Update app.py

app.py

@@ -10,8 +10,30 @@ You can also use 🔮DeepSeekMath📉 by cloning this space. Simply click here:
 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:"
 
+
+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}
+"""
+
+
 unimath1 = """Goal:
   
   X : UU
@@ -144,7 +166,8 @@ def parse_final_answer(answer):
 
 @spaces.GPU
 def solve_math_problem(question, additional_info, max_tokens):
-    prompt = f"User: {question}\n{additional_info}.\nAssistant:"
+    # Format the prompt using Lean4 structure
+    prompt = format_prompt(question, additional_info)
     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)
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -153,6 +176,7 @@ def solve_math_problem(question, additional_info, max_tokens):
     return full_answer, final_answer
 
 
+
 def main():
     iface = gr.Interface(        
         title="# 🙋🏻‍♂️Welcome to🌟Tonic's🔮Goedel Prover📉",