32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
|
||
|
|
|
||
|
|
from inference import build_prompt
|
||
|
|
|
||
|
|
MODEL_ID = 'EPFL-LARA/lean-standard-prover-qwen3-8b-grpo'
|
||
|
|
|
||
|
|
source = Path("Example.lean").read_text(encoding="utf-8")
|
||
|
|
prompt = build_prompt(source)
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
||
|
|
MODEL_ID, dtype=torch.bfloat16, device_map="auto", attn_implementation="sdpa",
|
||
|
|
).eval()
|
||
|
|
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
|
||
|
|
prompt_length = inputs["input_ids"].shape[1]
|
||
|
|
budget = min(2048, 32768 - prompt_length)
|
||
|
|
if budget <= 0:
|
||
|
|
raise ValueError("The prompt must leave room for a proof in the 32,768-token context.")
|
||
|
|
|
||
|
|
set_seed(0)
|
||
|
|
with torch.inference_mode():
|
||
|
|
output = model.generate(
|
||
|
|
**inputs, max_new_tokens=budget, do_sample=True,
|
||
|
|
temperature=0.3, top_p=0.95, top_k=0, min_p=0.0,
|
||
|
|
repetition_penalty=1.0, num_beams=1,
|
||
|
|
pad_token_id=tokenizer.pad_token_id,
|
||
|
|
)
|
||
|
|
response = tokenizer.decode(output[0, prompt_length:], skip_special_tokens=True)
|
||
|
|
print(response, end="")
|