Update app.py

app.py

@@ -15,9 +15,38 @@ def format_prompt(formal_statement, informal_prefix=""):
     return f"Complete the following Lean 4 code with explanatory comments preceding each line of code:\n\n```lean4\n{LEAN4_DEFAULT_HEADER}{informal_prefix}{formal_statement}"
 
 def extract_code(response):
-    """Extract code between lean4 code blocks"""
+    """Extract code between lean4 code blocks and the model's output"""
     try:
-        return re.search(r'```lean4\n(.*?)\n```', response, re.DOTALL).group(1)
+        # Extract the model's theorem and proof
+        output_match = re.search(r'theorem.*?end\s*```', response, re.DOTALL)
+        if output_match:
+            model_output = output_match.group(0)
+        else:
+            model_output = ""
+        
+        # Extract the initial code setup
+        setup_match = re.search(r'```lean4\n(.*?)============================', response, re.DOTALL)
+        if setup_match:
+            setup_code = setup_match.group(1)
+        else:
+            return "None"
+            
+        # Extract the goal
+        goal_match = re.search(r'============================\n\s*(.*?)(?:theorem|$)', response, re.DOTALL)
+        if goal_match:
+            goal = goal_match.group(1).strip()
+        else:
+            goal = ""
+            
+        # Combine all parts
+        full_code = (
+            f"{setup_code.strip()}\n"
+            f"============================\n"
+            f"{goal}\n\n"
+            f"{model_output.strip()}"
+        )
+        
+        return full_code
     except:
         return "None"