from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("shivamfet/slm-125m-instruct")
model = AutoModelForCausalLM.from_pretrained("shivamfet/slm-125m-instruct")
passage = ("The court affirmed the conviction but remanded for resentencing "
"because the trial judge improperly considered the defendant's silence.")
question = "Why did the court remand the case?"
messages = [
{"role": "system", "content": "Answer the question using only the passage."},
{"role": "user", "content": passage + "\n\n" + question},
]
prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
enc = tok(prompt, return_tensors="pt", return_token_type_ids=False)
eos = tok.convert_tokens_to_ids("<|eos|>")
out = model.generate(**enc, max_new_tokens=128, do_sample=False,
eos_token_id=eos, pad_token_id=eos)
print(tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip())