import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
model_name = "Qwen/Qwen2.5-1.5B-Instruct"
adapter = "beme08/qwen2.5-1.5b-riddle-lora"
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
model_name, torch_dtype=torch.float32, trust_remote_code=True
).to("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
model = PeftModel.from_pretrained(model, adapter)
def solve(riddle):
messages = [
{"role": "system", "content": "Solve the riddle. Respond with only the answer and no explanation."},
{"role": "user", "content": riddle},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=60, do_sample=False, pad_token_id=tokenizer.eos_token_id)
return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
print(solve("What gets wetter the more it dries?"))