adds mcp server, ssrmode false , refactor the global variables

app.py

@@ -5,13 +5,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import torch
 import json
 
-LEAN4_DEFAULT_HEADER = (
-    "import Mathlib\n"
-    "import Aesop\n\n"
-    "set_option maxHeartbeats 0\n\n"
-    "open BigOperators Real Nat Topology Rat\n"
-)
-
 title = "# 🙋🏻‍♂️Welcome to 🌟Tonic's 🌕💉👨🏻‍🔬Moonshot Math"
 
 description = """
@@ -39,27 +32,16 @@ joinus = """
 🌟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/qdfnvSPcqP) On 🤗Huggingface:[MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [MultiTonic](https://github.com/MultiTonic)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
 """
 
-# Build the initial system prompt
 SYSTEM_PROMPT = "You are an expert in mathematics and Lean 4."
 
-# Helper to build a Lean4 code block
-def build_formal_block(formal_statement, informal_prefix=""):
-    return (
-        f"{LEAN4_DEFAULT_HEADER}\n"
-        f"{informal_prefix}\n"
-        f"{formal_statement}"
-    )
 
-# Helper to extract the first Lean4 code block from text
-def extract_lean4_code(text):
-    code_block = re.search(r"```lean4(.*?)(```|$)", text, re.DOTALL)
-    if code_block:
-        code = code_block.group(1)
-        lines = [line for line in code.split('\n') if line.strip()]
-        return '\n'.join(lines)
-    return text.strip()
+LEAN4_DEFAULT_HEADER = (
+    "import Mathlib\n"
+    "import Aesop\n\n"
+    "set_option maxHeartbeats 0\n\n"
+    "open BigOperators Real Nat Topology Rat\n"
+)
 
-# Example problems
 unimath1 = """Goal:
   X : UU
   Y : UU
@@ -88,6 +70,21 @@ unimath4 = """Goal:
 
 additional_info_prompt = "/-Explain using mathematics-/\n"
 
+def build_formal_block(formal_statement, informal_prefix=""):
+    return (
+        f"{LEAN4_DEFAULT_HEADER}\n"
+        f"{informal_prefix}\n"
+        f"{formal_statement}"
+    )
+
+def extract_lean4_code(text):
+    code_block = re.search(r"```lean4(.*?)(```|$)", text, re.DOTALL)
+    if code_block:
+        code = code_block.group(1)
+        lines = [line for line in code.split('\n') if line.strip()]
+        return '\n'.join(lines)
+    return text.strip()
+
 examples = [
     [unimath1, additional_info_prompt, 1234],
     [unimath2, additional_info_prompt, 1234],
@@ -114,9 +111,7 @@ model_name = "AI-MO/Kimina-Prover-Distill-8B"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
 
-# Set generation config
 model.generation_config = GenerationConfig.from_pretrained(model_name)
-# Ensure pad_token_id is an integer, not a list
 if isinstance(model.generation_config.eos_token_id, list):
     model.generation_config.pad_token_id = model.generation_config.eos_token_id[0]
 else:
@@ -125,7 +120,6 @@ model.generation_config.do_sample = True
 model.generation_config.temperature = 0.6
 model.generation_config.top_p = 0.95
 
-# Initialize chat history with system prompt
 def init_chat(formal_statement, informal_prefix):
     user_prompt = (
         "Think about and solve the following problem step by step in Lean 4.\n"
@@ -137,15 +131,12 @@ def init_chat(formal_statement, informal_prefix):
         {"role": "user", "content": user_prompt}
     ]
 
-# Gradio chat handler
 @spaces.GPU
 def chat_handler(user_message, informal_prefix, max_tokens, chat_history):
-    # If chat_history is empty, initialize with system and first user message
     if not chat_history or len(chat_history) < 2:
         chat_history = init_chat(user_message, informal_prefix)
         display_history = [("user", user_message)]
     else:
-        # Append new user message
         chat_history.append({"role": "user", "content": user_message})
         display_history = []
         for msg in chat_history:
@@ -153,7 +144,6 @@ def chat_handler(user_message, informal_prefix, max_tokens, chat_history):
                 display_history.append(("user", msg["content"]))
             elif msg["role"] == "assistant":
                 display_history.append(("assistant", msg["content"]))
-    # Format prompt using chat template
     prompt = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
     attention_mask = torch.ones_like(input_ids)
@@ -166,14 +156,10 @@ def chat_handler(user_message, informal_prefix, max_tokens, chat_history):
         top_p=model.generation_config.top_p,
     )
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Extract only the new assistant message (after the prompt)
     new_response = result[len(prompt):].strip()
-    # Add assistant message to chat history
     chat_history.append({"role": "assistant", "content": new_response})
     display_history.append(("assistant", new_response))
-    # Extract Lean4 code
     code = extract_lean4_code(new_response)
-    # Prepare output
     output_data = {
         "model_input": prompt,
         "model_output": result,
@@ -184,7 +170,6 @@ def chat_handler(user_message, informal_prefix, max_tokens, chat_history):
 
 def main():
     with gr.Blocks() as demo:
-        # Title and Model Description
         gr.Markdown("""# 🙋🏻‍♂️Welcome to 🌟Tonic's 🌕💉👨🏻‍🔬Moonshot Math""")
         with gr.Row():
             with gr.Column():
@@ -213,4 +198,4 @@ def main():
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main(ssr_mode=False, mcp_server=True)