Spaces:
Running
on
Zero
Running
on
Zero
| import spaces | |
| import re | |
| import gradio as gr | |
| 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 = """ | |
| **🌕💉👨🏻🔬AI-MO/Kimina-Prover-Distill-8B is a theorem proving model developed by Project Numina and Kimi teams, focusing on competition style problem solving capabilities in Lean 4. It is a distillation of AI-MO/Kimina-Prover-72B, a model trained via large scale reinforcement learning. It achieves 77.86% accuracy with Pass@32 on MiniF2F-test.\ | |
| - [Kimina-Prover-Preview GitHub](https://github.com/MoonshotAI/Kimina-Prover-Preview)\ | |
| - [Hugging Face: AI-MO/Kimina-Prover-72B](https://huggingface.co/AI-MO/Kimina-Prover-72B)\ | |
| - [Kimina Prover blog](https://huggingface.co/blog/AI-MO/kimina-prover)\ | |
| - [unimath dataset](https://huggingface.co/datasets/introspector/unimath)\ | |
| """ | |
| citation = """> **Citation:** | |
| > ``` | |
| > @article{kimina_prover_2025, | |
| > title = {Kimina-Prover Preview: Towards Large Formal Reasoning Models with Reinforcement Learning}, | |
| > author = {Wang, Haiming and Unsal, Mert and ...}, | |
| > year = {2025}, | |
| > url = {http://arxiv.org/abs/2504.11354}, | |
| > } | |
| > ``` | |
| """ | |
| joinus = """ | |
| ## Join us : | |
| 🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [](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() | |
| # Example problems | |
| unimath1 = """Goal: | |
| X : UU | |
| Y : UU | |
| P : UU | |
| xp : (X → P) → P | |
| yp : (Y → P) → P | |
| X0 : X × Y → P | |
| x : X | |
| ============================ | |
| (Y → P)""" | |
| unimath2 = """Goal: | |
| R : ring M : module R | |
| ============================ | |
| (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))""" | |
| 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))""" | |
| additional_info_prompt = "/-Explain using mathematics-/\n" | |
| examples = [ | |
| [unimath1, additional_info_prompt, 2500], | |
| [unimath2, additional_info_prompt, 2500], | |
| [unimath3, additional_info_prompt, 2500], | |
| [unimath4, additional_info_prompt, 2500], | |
| # New examples | |
| [ | |
| '''import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n/-- Let $a_1, a_2,\cdots, a_n$ be real constants, $x$ a real variable, and $f(x)=\\cos(a_1+x)+\\frac{1}{2}\\cos(a_2+x)+\\frac{1}{4}\\cos(a_3+x)+\\cdots+\\frac{1}{2^{n-1}}\\cos(a_n+x).$ Given that $f(x_1)=f(x_2)=0,$ prove that $x_2-x_1=m\\pi$ for some integer $m.$-/\ntheorem imo_1969_p2 (m n : \\R) (k : \\N) (a : \\N \\rightarrow \\R) (y : \\R \\rightarrow \\R) (h₀ : 0 < k)\n(h₁ : \\forall x, y x = \\sum i in Finset.range k, Real.cos (a i + x) / 2 ^ i) (h₂ : y m = 0)\n(h₃ : y n = 0) : \\exists t : \\Z, m - n = t * Real.pi := by''', | |
| "/-- Let $a_1, a_2,\\cdots, a_n$ be real constants, $x$ a real variable, and $f(x)=\\cos(a_1+x)+\\frac{1}{2}\\cos(a_2+x)+\\frac{1}{4}\\cos(a_3+x)+\\cdots+\\frac{1}{2^{n-1}}\\cos(a_n+x).$ Given that $f(x_1)=f(x_2)=0,$ prove that $x_2-x_1=m\\pi$ for some integer $m.$-/", | |
| 2500 | |
| ], | |
| [ | |
| '''import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n/-- Suppose that $h(x)=f^{-1}(x)$. If $h(2)=10$, $h(10)=1$ and $h(1)=2$, what is $f(f(10))$? Show that it is 1.-/\ntheorem mathd_algebra_209 (σ : Equiv \\R \\R) (h₀ : σ.2 2 = 10) (h₁ : σ.2 10 = 1) (h₂ : σ.2 1 = 2) :\nσ.1 (σ.1 10) = 1 := by''', | |
| "/-- Suppose that $h(x)=f^{-1}(x)$. If $h(2)=10$, $h(10)=1$ and $h(1)=2$, what is $f(f(10))$? Show that it is 1.-/", | |
| 2500 | |
| ], | |
| [ | |
| '''import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n/-- At which point do the lines $s=9-2t$ and $t=3s+1$ intersect? Give your answer as an ordered pair in the form $(s, t).$ Show that it is (1,4).-//\ntheorem mathd_algebra_44 (s t : \\R) (h₀ : s = 9 - 2 * t) (h₁ : t = 3 * s + 1) : s = 1 \\wedge t = 4 := by''', | |
| "/-- At which point do the lines $s=9-2t$ and $t=3s+1$ intersect? Give your answer as an ordered pair in the form $(s, t).$ Show that it is (1,4).-/", | |
| 2500 | |
| ], | |
| ] | |
| 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: | |
| model.generation_config.pad_token_id = model.generation_config.eos_token_id | |
| 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" | |
| "# Problem: Provide a formal proof for the following statement.\n" | |
| f"# Formal statement:\n```lean4\n{build_formal_block(formal_statement, informal_prefix)}\n```\n" | |
| ) | |
| return [ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": user_prompt} | |
| ] | |
| # Gradio chat handler | |
| 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: | |
| if msg["role"] == "user": | |
| 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) | |
| outputs = model.generate( | |
| input_ids, | |
| attention_mask=attention_mask, | |
| max_length=max_tokens + input_ids.shape[1], | |
| pad_token_id=model.generation_config.pad_token_id, | |
| temperature=model.generation_config.temperature, | |
| 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, | |
| "lean4_code": code, | |
| "chat_history": chat_history | |
| } | |
| return display_history, json.dumps(output_data, indent=2), code, 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(): | |
| gr.Markdown(description) | |
| with gr.Column(): | |
| gr.Markdown(joinus) | |
| with gr.Row(): | |
| with gr.Column(): | |
| user_input = gr.Textbox(label="👨🏻💻Your message or formal statement", lines=4) | |
| informal = gr.Textbox(value=additional_info_prompt, label="💁🏻♂️Optional informal prefix") | |
| max_tokens = gr.Slider(minimum=150, maximum=4096, value=2500, label="🪙Max Tokens") | |
| submit = gr.Button("Send") | |
| with gr.Column(): | |
| chat = gr.Chatbot(label="🌕💉👨🏻🔬Kimina Prover 8B") | |
| with gr.Accordion("Complete Output", open=False): | |
| json_out = gr.JSON(label="Full Output") | |
| code_out = gr.Code(label="Extracted Lean4 Code", language="python") | |
| state = gr.State([]) | |
| submit.click(chat_handler, [user_input, informal, max_tokens, state], [chat, json_out, code_out, state]) | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[user_input, informal, max_tokens], | |
| label="🤦🏻♂️Example Problems" | |
| ) | |
| gr.Markdown(citation) | |
| demo.launch() | |
| if __name__ == "__main__": | |
| main() | |