Update app.py

app.py

@@ -3,39 +3,26 @@ import re
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
+import json
 
-title = """# 🙋🏻‍♂️Welcome to🌟Tonic's🔮DeepSeekMath📉
-You can build with this endpoint using🔮DeepSeekMath📉 available here : [deepseek-ai/deepseek-math-7b-instruct](https://huggingface.co/deepseek-ai/deepseek-math-7b-instruct). 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 🔮DeepSeekMath📉 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 🤗
-"""
-
-
-LEAN4_DEFAULT_HEADER = """import Mathlib
-import Aesop
-
-set_option maxHeartbeats 0
-
-open BigOperators Real Nat Topology Rat
-
-"""
-
-
-additional_info_prompt = "Explain the above using mathematics, print entire answer in latex format:"
+LEAN4_DEFAULT_HEADER = "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n"
 
+title = """# 🙋🏻‍♂️Welcome to🌟Tonic's🔮Goedel Prover📉
+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)."""
 
-def format_prompt(question, additional_info):
-    return f"""Complete the following Lean 4 code with explanatory comments preceding each line of code:
-
-```lean4
-{LEAN4_DEFAULT_HEADER}
-{question}
-{additional_info}
-"""
+def format_prompt(formal_statement, informal_prefix=""):
+    """Format the input according to the Lean4 structure"""
+    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"""
+    try:
+        return re.search(r'```lean4\n(.*?)\n```', response, re.DOTALL).group(1)
+    except:
+        return "None"
 
+# Example problems
 unimath1 = """Goal:
-  
   X : UU
   Y : UU
   P : UU
@@ -44,95 +31,24 @@ unimath1 = """Goal:
   X0 : X × Y → P
   x : X
   ============================
-   (Y → P)
-
-
-DEBUG:Going to execute:
-PTRDEBUGTAC<coq-core.plugins.ltac::intro@1> $1
-DEBUG LTAC Evaluated term: yp
-
-TcDebug (0) > NONE??? LtacNotationCall (no location),/mnt/data1/2024/01/05/UniMath/UniMath/Foundations/PartA.v:234
-"""
-# source : unimath/unimath/batch2/data08
+   (Y → P)"""
 
 unimath2 = """Goal:
     R : ring  M : module R
   ============================
-   (islinear (idfun M))
-
-
-DEBUG:Going to execute:
-PTRDEBUGTACapply pathsinv0; trivial
-Level 0: Backtrace:
-
-Proof is not complete.
-Level 0: Backtrace:
-
-Proof is not complete.
-
-TcDebug (0) > NONE??? LtacNotationCall (no location),/mnt/data1/2024/01/05/UniMath/UniMath/MoreFoundations/Tactics.veasy:19, Tactics  (UniMath.MoreFoundations),/mnt/data1/2024/01/05/UniMath/UniMath/Algebra/Modules/Examples.v:27
-
-"""
-# source : unimath/unimath/batch2/data_22/BATCH122007
+   (islinear (idfun M))"""
 
 unimath3 = """Goal:
     X : UU  i : nat  b : hProptoType (i < S i)  x : Vector X (S i)  r : i = i
   ============================
-   (pr1 lastelement = pr1 (i,, b))
-
-
-DEBUG:Going to execute:
-PTRDEBUGTACsimpl
-DEBUG LTAC Evaluated term: isinjstntonat
-
-TcDebug (0) > /mnt/data1/2024/01/05/UniMath/UniMath/Combinatorics/FiniteSequences.v:114
-"""
-# source : unimath/unimath/batch2/data_12/BATCH112026
+   (pr1 lastelement = pr1 (i,, b))"""
 
 unimath4 = """Goal:
     X : dcpo  CX : continuous_dcpo_struct X  x : pr1hSet X  y : pr1hSet X
   ============================
-   (x ⊑ y
-    ≃ (∀ i : approximating_family CX x, approximating_family CX x i ⊑ y))
-
-
-DEBUG:Going to execute:
-PTRDEBUGTACsimple refine (p _ _ _) ||
-             simple refine (p _ _ _ _) ||
-               simple refine (p _ _ _ _ _) ||
-                 simple refine (p _ _ _ _ _ _) ||
-                   simple refine (p _ _ _ _ _ _ _) ||
-                     simple refine (p _ _ _ _ _ _ _ _) ||
-                       simple refine (p _ _ _ _ _ _ _ _ _) ||
-                         simple refine (p _ _ _ _ _ _ _ _ _ _) ||
-                           simple refine (p _ _ _ _ _ _ _ _ _ _ _) ||
-                             simple refine (p _ _ _ _ _ _ _ _ _ _ _ _) ||
-                               simple refine (p _ _ _ _ _ _ _ _ _ _ _ _ _) ||
-                                 simple refine
-                                  (p _ _ _ _ _ _ _ _ _ _ _ _ _ _) || simple
-                                  refine (p _ _ _ _ _ _ _ _ _ _ _ _ _ _ _)
-Level 0: Backtrace:
-
-In environment
-X : dcpo
-CX : continuous_dcpo_struct X
-x, y : X
-The term "weqimplimpl ?f ?g" has type "isaprop ?X → isaprop ?Y → ?X ≃ ?Y"
-while it is expected to have type
- "x ⊑ y ≃ (∀ i : approximating_family CX x, approximating_family CX x i ⊑ y)".
-Level 0: Backtrace:
-
-In environment
-X : dcpo
-CX : continuous_dcpo_struct X
-x, y : X
-The term "weqimplimpl ?f ?g" has type "isaprop ?X → isaprop ?Y → ?X ≃ ?Y"
-while it is expected to have type
- "x ⊑ y ≃ (∀ i : approximating_family CX x, approximating_family CX x i ⊑ y)".
-
-TcDebug (0) > NONE??? LtacNotationCall (no location),/mnt/data1/2024/01/05/UniMath/UniMath/Foundations/Init.vsimple_rapply:174, Init  (UniMath.Foundations),/mnt/data1/2024/01/05/UniMath/UniMath/Foundations/Init.v??? LtacNotationCall:189, Init  (UniMath.Foundations),/mnt/data1/2024/01/05/UniMath/UniMath/OrderTheory/DCPOs/Basis/Continuous.v:166
-"""
-# source : unimath/unimath/batch2/data_42/BATCH142042
+   (x ⊑ y ≃ (∀ i : approximating_family CX x, approximating_family CX x i ⊑ y))"""
+
+additional_info_prompt = "/-Explain using mathematics-/"
 
 examples = [
     [unimath1, additional_info_prompt, 1200],
@@ -144,8 +60,6 @@ examples = [
 model_name = "Goedel-LM/Goedel-Prover-SFT"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
-# model.generation_config = GenerationConfig.from_pretrained(model_name)
-# model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
 # Set generation config
 model.generation_config = GenerationConfig.from_pretrained(model_name)
@@ -153,29 +67,33 @@ model.generation_config.pad_token_id = model.generation_config.eos_token_id
 model.generation_config.bos_token_id = 100000
 model.generation_config.eos_token_id = 100001
 
-
-def parse_full_answer(answer):
-    """Parses the assistant's answer, excluding any text before 'Assistant :'."""
-    match = re.search(r"Assistant\s*:\s*(.*)", answer, re.DOTALL)
-    return match.group(1).strip() if match else "No assistant answer found."
-
-def parse_final_answer(answer):
-    """Extracts the final answer enclosed within \boxed{}."""
-    match = re.search(r"\\boxed\{([^}]+)\}", answer)
-    return match.group(1) if match else "No final answer found."
-
 @spaces.GPU
-def solve_math_problem(question, additional_info, max_tokens):
+def solve_math_problem(question, informal_prefix, max_tokens):
     # Format the prompt using Lean4 structure
-    prompt = format_prompt(question, additional_info)
+    prompt = format_prompt(question, informal_prefix)
+    
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
-    outputs = model.generate(input_ids, max_length=max_tokens + input_ids.shape[1], pad_token_id=model.generation_config.pad_token_id)
+    outputs = model.generate(
+        input_ids,
+        max_length=max_tokens + input_ids.shape[1],
+        pad_token_id=model.generation_config.pad_token_id,
+        temperature=1.0,
+        top_p=0.95,
+    )
+    
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    full_answer = parse_full_answer(result)
-    final_answer = parse_final_answer(result)
-    return full_answer, final_answer
-
-
+    
+    # Extract the full code from the response
+    full_code = extract_code(prompt + result)
+    
+    # Create output dictionary similar to reference code
+    output_data = {
+        "model_input": prompt,
+        "model_output": result,
+        "full_code": full_code
+    }
+    
+    return json.dumps(output_data, indent=2), full_code
 
 def main():
     iface = gr.Interface(        
@@ -184,18 +102,15 @@ def main():
 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 🤗
 """,
-
-
-        
         fn=solve_math_problem,
         outputs=[
-            gr.Code(label="🔮TonicsMathAssistant📉", interactive=False),
-            gr.Textbox(label="Final Answer")
+            gr.JSON(label="Full Output"),
+            gr.Code(label="Extracted Lean4 Code", language="lean")
         ],
         inputs=[
-            gr.Textbox(label="🤔Enter your math problem", lines=7),
-            gr.Textbox(value=additional_info_prompt, label="🪜Optional train-of-thought"),
-            gr.Slider(minimum=150, maximum=1200, value=650, label="🪙Max Tokens")
+            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")
         ],
         examples=examples
     )
@@ -203,4 +118,4 @@ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder
     iface.launch()
 
 if __name__ == "__main__":
-    main()
+    main()