import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
BASE = "Qwen/Qwen3-4B-Instruct-2507"
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
tok = AutoTokenizer.from_pretrained(BASE)
m = AutoModelForCausalLM.from_pretrained(BASE, quantization_config=bnb, device_map="auto")
m = PeftModel.from_pretrained(m, "William2390401/aime-gen-qwen3-4b-lora-v3")
msgs = [{"role": "user", "content":
"Write an AIME-style problem. Difficulty: late (problems 11-15). Topic: number theory."}]
text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
out = m.generate(**tok(text, return_tensors="pt").to(m.device),
max_new_tokens=1600, do_sample=True, temperature=0.8, top_p=0.95)
print(tok.decode(out[0], skip_special_tokens=True))