import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("jonam-ai/legal-slm-125m-raft")
model = AutoModelForCausalLM.from_pretrained("jonam-ai/legal-slm-125m-raft", torch_dtype=torch.float32)
sid = tok.convert_tokens_to_ids
system = ("You are a legal and financial assistant. Use the numbered context documents to "
"answer the question. Quote the text you rely on, then give the final answer.")
context = ("Context:\n[1] The Company entered into a five-year lease at an annual rent of $2.4 million.\n"
"[2] The board declared a quarterly dividend of $0.15 per share.")
question = "What is the annual rent for the lease?"
ids = (tok("<|bos|>", add_special_tokens=False)["input_ids"]
+ [sid("<|system|>")] + tok(system, add_special_tokens=False)["input_ids"]
+ [sid("<|user|>")] + tok(f"{context}\n\nQuestion: {question}", add_special_tokens=False)["input_ids"]
+ [sid("<|assistant|>")])
out = model.generate(torch.tensor([ids]), max_new_tokens=160, do_sample=True, temperature=0.5,
eos_token_id=sid("<|eos|>"), pad_token_id=sid("<|pad|>"))
print(tok.decode(out[0][len(ids):], skip_special_tokens=True))