Views
No views yet
1from vllm import LLM, SamplingParams
2from transformers import AutoTokenizer
3
4model_name = "Stepfun/Stepfun-Prover-Preview-32B"
5model = LLM(
6 model=model_name,
7 tensor_parallel_size=8,
8 )
9tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
10
11formal_problem = """
12import Mathlib
13
14theorem test_theorem (x y z : ℝ) (hx : 0 < x) (hy : 0 < y) (hz : 0 < z) :
15 (x^2 - z^2) / (y + z) + (y^2 - x^2) / (z + x) + (z^2 - y^2) / (x + y) ≥ 0 := by
16""".strip()
17
18system_prompt = "You will be given an unsolved Lean 4 problem. Think carefully and work towards a solution. At any point, you may use the Lean 4 REPL to check your progress by enclosing your partial solution between <sketch> and </sketch>. The REPL feedback will be provided between <REPL> and </REPL>. Continue this process as needed until you arrive at a complete and correct solution."
19
20user_prompt = f"```lean4\n{formal_problem}\n```"
21
22dialog = [
23 {"role": "system", "content": system_prompt},
24 {"role": "user", "content": user_prompt}
25]
26
27prompt = tokenizer.apply_chat_template(dialog, tokenize=False, add_generation_prompt=True)
28
29sampling_params = SamplingParams(
30 temperature=0.999,
31 top_p=0.95,
32 top_k=-1,
33 max_tokens=16384,
34 stop_token_ids=[151643, 151666], # <|end▁of▁sentence|>, </sketch>
35 include_stop_str_in_output=True,
36)
37
38output = model.generate(prompt, sampling_params=sampling_params)
39output_text = output[0].outputs[0].text
40print(output_text)