Update app.py

app.py

@@ -267,7 +267,106 @@ def calculate_code_metrics(reference, generated):
             'char_overlap': 0
         }
 
-def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p, num_sequences, reference_code):
+def format_python_code(code):
+    """Format and clean generated code to be proper Python syntax with indentation"""
+    if not code or code.startswith('#'):
+        return code
+
+    try:
+        # Basic Python syntax corrections (convert common C++/other language patterns to Python)
+        corrections = [
+            # Data types
+            ('int ', ''),  # Remove C-style type declarations
+            ('string ', ''),  # Remove string type declarations
+            ('bool ', ''),  # Remove bool type declarations
+            ('float ', ''),  # Remove float type declarations
+            ('void ', ''),  # Remove void declarations
+            ('char ', ''),  # Remove char type declarations
+            ('double ', ''),  # Remove double type declarations
+
+            # Function declarations
+            ('function ', 'def '),  # Convert function to def
+            ('public ', ''),  # Remove public
+            ('private ', ''),  # Remove private
+            ('static ', ''),  # Remove static
+            ('const ', ''),  # Remove const
+
+            # Control structures
+            ('elseif', 'elif'),  # Convert elseif to elif
+            ('else if', 'elif'),  # Convert else if to elif
+
+            # Loops
+            ('for(int ', 'for '),  # Remove int from for loops
+            ('for(string ', 'for '),  # Remove string from for loops
+            ('for(char ', 'for '),  # Remove char from for loops
+
+            # Common patterns
+            ('System.out.println', 'print'),  # Convert Java print to Python
+            ('console.log', 'print'),  # Convert JS print to Python
+            ('cout <<', 'print'),  # Convert C++ cout to Python
+            ('printf', 'print'),  # Convert C printf to Python
+            ('std::', ''),  # Remove C++ std namespace
+            ('->', '.'),  # Convert arrow operator to dot
+            ('::', '.'),  # Convert scope resolution to dot
+
+            # Braces and brackets
+            ('{', ':'),  # Convert opening braces to colons
+            ('}', ''),  # Remove closing braces
+
+            # Semicolons (remove them)
+            ('; ', '\n'),  # Convert semicolons to newlines
+            (';', '\n'),  # Convert semicolons to newlines
+        ]
+
+        for old, new in corrections:
+            code = code.replace(old, new)
+
+        # Handle function calls and parentheses
+        # Convert function(param) to function(param): for function definitions
+        import re
+        # Pattern to find function definitions like "def func(param)" and add colon
+        code = re.sub(r'(def\s+\w+\([^)]*\))(?!\s*:)', r'\1:', code)
+
+        # Split into lines and add proper indentation
+        lines = code.split('\n')
+        formatted_lines = []
+        indent_level = 0
+        indent_size = 4  # 4 spaces per indent level
+
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+
+            # Decrease indent for dedent keywords
+            if any(line.startswith(keyword) for keyword in ['else:', 'elif ', 'except:', 'finally:', 'return', 'break', 'continue', 'pass']):
+                indent_level = max(0, indent_level - 1)
+
+            # Add indentation
+            indented_line = ' ' * (indent_level * indent_size) + line
+            formatted_lines.append(indented_line)
+
+            # Increase indent after indent keywords
+            if ':' in line and not line.startswith(('else:', 'elif ', 'except:', 'finally:')):
+                if any(keyword in line for keyword in ['if ', 'for ', 'while ', 'def ', 'class ', 'try:', 'except ', 'with ']):
+                    indent_level += 1
+
+        # Join lines back
+        formatted_code = '\n'.join(formatted_lines)
+
+        # Additional cleanup
+        formatted_code = re.sub(r'\s+', ' ', formatted_code)  # Remove multiple spaces
+        formatted_code = re.sub(r'\n\s*\n\s*\n', '\n\n', formatted_code)  # Remove excessive newlines
+
+        # Ensure proper Python syntax
+        # Add colons where missing for control structures
+        formatted_code = re.sub(r'\b(if|for|while|def|class|try|except|with)\s+([^:]+)(?!\s*:)', r'\1 \2:', formatted_code)
+
+        return formatted_code
+
+    except Exception as e:
+        # If formatting fails, return original code with a note
+        return f"# Formatting failed: {str(e)}\n{code}"def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p, num_sequences, reference_code):
     """Generate code from pseudo-code using loaded model"""
     global loaded_model, loaded_tokenizer, generation_history
     
@@ -280,8 +379,8 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
     try:
         start_time = time.time()
         
-        # Format input
-        prompt = f"<PSEUDO> {pseudo_code.strip()} <SEP> <CODE>"
+        # Format input with Python-specific instructions
+        prompt = f"<PSEUDO> {pseudo_code.strip()} <SEP> Generate Python code with proper indentation and syntax: <CODE>"
         
         # Tokenize with error handling
         device = next(loaded_model.parameters()).device
@@ -500,8 +599,11 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
                     for token in special_tokens:
                         code = code.replace(token, '')
                     
-                    # Clean up extra whitespace
-                    code = ' '.join(code.split())
+                    # Clean up extra whitespace but preserve some structure
+                    code = code.replace('\n\n\n', '\n\n')  # Reduce excessive newlines
+                    
+                    # Format as proper Python code
+                    code = format_python_code(code)
                     
                     # Ensure we have some content
                     if not code or code.isspace():