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="")