Update app.py

app.py

@@ -28,36 +28,33 @@ def format_prompt(formal_statement, informal_prefix=""):
 def extract_code(response):
     """Extract code between lean4 code blocks and the model's output"""
     try:
-        # 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:
+        # Find all content between ```lean4 and ``` tags
+        matches = re.findall(r'```lean4(.*?)```', response, re.DOTALL)
+        if not matches:
             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
+        # Get the last match which should contain the complete code including model output
+        full_content = matches[-1].strip()
+        
+        # Clean up any duplicate headers or content
+        # Split the content into lines
+        lines = full_content.split('\n')
+        cleaned_lines = []
+        seen_headers = set()
+        
+        for line in lines:
+            # Skip empty lines at the start
+            if not cleaned_lines and not line.strip():
+                continue
+            # Skip duplicate headers
+            if "import Mathlib" in line or "import Aesop" in line or "set_option" in line:
+                if line in seen_headers:
+                    continue
+                seen_headers.add(line)
+            cleaned_lines.append(line)
+        
+        # Join the lines back together
+        return '\n'.join(cleaned_lines)
     except:
         return "None"
 
@@ -146,7 +143,7 @@ def main():
         title="# 🙋🏻‍♂️Welcome to🌟Tonic's🔮Goedel Prover📉",
         description="""You can build with this endpoint using🔮Goedel-Prover-SFT📉 available here : [Goedel-LM/Goedel-Prover-SFT](https://huggingface.co/Goedel-LM/Goedel-Prover-SFT). We're using 🤖[introspector/unimath](https://huggingface.co/datasets/introspector/unimath) for cool examples, check it out below ! The demo is still a work in progress and we're looking forward to build downstream tasks that showcase outstanding mathematical reasoning. Have any ideas ? join us below !
 You can also use 🔮Goedel Prover📉 by cloning this space. Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/Math?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
-Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) Math with [introspector](https://huggingface.co/introspector) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [SciTonic](https://github.com/Tonic-AI/scitonic)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
+Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [Join us on Discord](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) Math with [introspector](https://huggingface.co/introspector) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [SciTonic](https://github.com/Tonic-AI/scitonic)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
 """,
         fn=solve_math_problem,
         outputs=[
@@ -156,7 +153,7 @@ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder
         inputs=[
             gr.Textbox(label="🤔Enter your Lean4 formal statement", lines=7),
             gr.Textbox(value=additional_info_prompt, label="🪜Optional informal prefix"),
-            gr.Slider(minimum=150, maximum=2048, value=650, label="🪙Max Tokens")
+            gr.Slider(minimum=150, maximum=4086, value=2500, label="🪙Max Tokens")
         ],
         examples=examples
     )