from vllm import LLM, SamplingParamsfrom transformers import AutoTokenizer def get_formal_statement_prompt(informal_problem: str, header: str = "import Mathlib\n") -> str: prompt = "Please autoformalize the following problem in Lean 4 with a header. Use the following theorem names: my_favorite_theorem.\n\n" prompt += informal_problem prompt += f"\n\nYour code should start with:\n```Lean4\n{header}\n```\n" return prompt MODEL_DIR = "stepfun-ai/StepFun-Formalizer-32B" if __name__ == "__main__": system_prompt = "You are an expert in mathematics and Lean 4." informal_problem = "The real numbers $x, y, z$ satisfy $0 \\leq x \\leq y \\leq z \\leq 4$. If their squares form an arithmetic progression with common difference 2, determine the minimum possible value of $|x-y|+|y-z|$.\n Prove that the answer is: 4-2\\sqrt{3}" header = "import Mathlib\n\nopen Real\n" user_prompt = get_formal_statement_prompt(informal_problem, header) dialog = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ] tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR) prompt = tokenizer.apply_chat_template(dialog, tokenize=False, add_generation_prompt=True) + "<think>" print(f"prompt: {prompt}") model = LLM( MODEL_DIR, tensor_parallel_size=4 # 8 for 32B, 4 for 7B ) sampling_params = SamplingParams( temperature=0.6, top_p=0.95, max_tokens=16384, n=1 ) responses = model.generate(prompt, sampling_params) print(f"response: {responses[0].outputs[0].text}")