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 👻 [![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() # 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 @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: 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()