Update app.py

app.py

@@ -273,118 +273,76 @@ def format_python_code(code):
         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 to Python indentation
-            ('{', ':\n'),  # Convert opening braces to colons with newline
-            ('}', ''),  # Remove closing braces
-
-            # Semicolons (convert to newlines)
-            ('; ', '\n'),  # Convert semicolons to newlines
-            (';', '\n'),  # Convert semicolons to newlines
-
-            # Ternary operators (C++ style to Python)
-            (' ? ', ' if '),  # Convert ternary ? to if
-            (' : ', ' else '),  # Convert ternary : to else
-
-            # Comparison operators
-            ('< SEP', '<'),  # Fix common typo
-            ('< SEp', '<'),  # Fix common typo
-            ('<SEP', '<'),  # Fix common typo
-            ('<SEp', '<'),  # Fix common typo
-        ]
-
-        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')
+        
+        # Remove special tokens and artifacts first
+        code = re.sub(r'<[^>]*>', '', code)  # Remove all <TOKEN> patterns
+        code = code.replace('<TR>', '').strip()  # Remove <TR> specifically
+        
+        # Basic cleanup and conversion to Python
+        # Convert C++ function declarations to Python
+        code = re.sub(r'\b(?:bool|int|void|string|float|char|double)\s+(\w+)\s*\(([^)]*)\)\s*\{', r'def \1(\2):', code)
+        
+        # Clean up parameter types in function signatures
+        code = re.sub(r'\(\s*(?:int|bool|string|float|char|double)\s+(\w+)\s*\)', r'(\1)', code)
+        code = re.sub(r',\s*(?:int|bool|string|float|char|double)\s+(\w+)', r', \1', code)
+        
+        # Replace braces with proper Python structure
+        code = code.replace('{', ':')
+        code = code.replace('}', '')
+        
+        # Remove semicolons
+        code = code.replace(';', '')
+        
+        # Fix return statements
+        code = re.sub(r'return\s+true\b', 'return True', code)
+        code = re.sub(r'return\s+false\b', 'return False', code)
+        
+        # Fix control structures
+        code = re.sub(r'\bif\s*\(([^)]+)\)', r'if \1:', code)
+        code = re.sub(r'\belse\s*:', r'else:', code)
+        code = re.sub(r'\belse\s+', r'else:\n    ', code)
+        
+        # Split into lines for indentation
+        lines = [line.strip() for line in code.split('\n') if line.strip()]
+        
+        # Add proper indentation
         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']):
+            # Handle dedent
+            if line.startswith('else:') or line.startswith('elif'):
                 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)
-
-        # Fix ternary operators that became malformed
-        # Convert "return a if condition else b" to proper Python ternary
-        formatted_code = re.sub(r'return\s+(.+?)\s+if\s+(.+?)\s+else\s+(.+)', r'return \1 if \2 else \3', formatted_code)
-
-        # Ensure functions have proper structure
-        # If a line starts with "def" and doesn't end with ":", add it
-        formatted_code = re.sub(r'(def\s+\w+\([^)]*)\s*$', r'\1:', formatted_code, flags=re.MULTILINE)
-
-        return formatted_code
+            if indent_level > 0:
+                formatted_line = '    ' * indent_level + line
+            else:
+                formatted_line = line
+            
+            formatted_lines.append(formatted_line)
+            
+            # Handle indent after colon
+            if line.endswith(':'):
+                indent_level += 1
+        
+        # Join lines
+        result = '\n'.join(formatted_lines)
+        
+        # Final cleanup
+        result = re.sub(r'\n\s*\n+', '\n', result)  # Remove empty lines
+        
+        # Ensure we have something useful
+        if not result.strip() or 'def ' not in result:
+            # Create a basic function if parsing failed
+            result = f"def generated_function():\n    # Model output: {code[:50]}...\n    return None"
+        
+        return result
 
     except Exception as e:
-        # If formatting fails, return original code with a note
-        return f"# Formatting failed: {str(e)}\n{code}"
+        # If formatting fails, return a basic structure with the original
+        return f"def generated_function():\n    # Formatting error: {str(e)}\n    # Original: {code[:100]}...\n    return None"
 
 
 def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p, num_sequences, reference_code):
@@ -401,7 +359,7 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
         start_time = time.time()
         
         # Format input with Python-specific instructions
-        prompt = f"<PSEUDO> {pseudo_code.strip()} <SEP> Generate proper Python code with correct syntax, indentation, and separate lines. Use def for functions, if/else for conditionals, for/while for loops, and proper Python syntax only. No C++ or other languages. <CODE>"
+        prompt = f"<PSEUDO> {pseudo_code.strip()} <SEP> Write a Python function to {pseudo_code.strip()}. Use proper Python syntax with def, return statements, if/else conditions, and proper indentation. Example: def check_even_odd(number): if number % 2 == 0: return 'even' else: return 'odd' <CODE>"
         
         # Tokenize with error handling
         device = next(loaded_model.parameters()).device
@@ -623,13 +581,15 @@ def generate_code_from_pseudo(pseudo_code, max_length, temperature, top_k, top_p
                     # 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)
-                    
                     # For debugging: include raw generated code
                     raw_code = generated.strip()
-                    if not code.startswith('#'):
-                        code = f"# Raw generated code:\n# {raw_code}\n\n# Formatted Python code:\n{code}"
+                    formatted_code = format_python_code(code)
+                    
+                    # Show both raw and formatted for transparency
+                    if not formatted_code.startswith('#'):
+                        code = f"# Model Generated (Raw):\n# {raw_code[:100]}...\n\n# Formatted Python Code:\n{formatted_code}"
+                    else:
+                        code = formatted_code
                     
                     # Ensure we have some content
                     if not code or code.isspace():