import torch
from transformers import AutoModelForCausalLM, AutoProcessor
from peft import PeftModel
BASE = "unsloth/gemma-4-E2B-it"
ADAPTER = "xbruce22/gemma-4-e2b-reasoning-lora"
device = "cuda" if torch.cuda.is_available() else (
"xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else "cpu")
dtype = torch.float32 if device == "cpu" else torch.bfloat16
base = AutoModelForCausalLM.from_pretrained(BASE, dtype=dtype).to(device)
model = PeftModel.from_pretrained(base, ADAPTER)
model = model.merge_and_unload()
model.eval()
processor = AutoProcessor.from_pretrained(BASE)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Write DFS in python, keep short."},
]
text = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
inputs = processor(text=[text], return_tensors="pt").to(device)
for k in list(inputs):
if "token_type" in k or "pixel" in k or "audio" in k:
inputs.pop(k)
with torch.inference_mode():
out = model.generate(
**inputs, max_new_tokens=1024, do_sample=True,
temperature=1.0, top_p=0.95, top_k=64,
pad_token_id=processor.tokenizer.pad_token_id)
gen = out[0][inputs["input_ids"].shape[1]:]
print(processor.decode(gen, skip_special_tokens=True))