File size: 997 Bytes
45b6ae4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL = "Project-Prism/Orion-Flagship-Mini-SFT"
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
MODEL,
trust_remote_code=True,
torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
device_map="auto" if torch.cuda.is_available() else None,
)
model.eval()
messages = [
{"role": "system", "content": "You are Orion, a helpful AI assistant."},
{"role": "user", "content": "What is 7 x 10?"},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.inference_mode():
out = model.generate(**inputs, max_new_tokens=64, temperature=0.7, top_p=0.95, do_sample=True, use_cache=False)
answer_ids = out[0, inputs["input_ids"].shape[1]:]
print(tokenizer.decode(answer_ids, skip_special_tokens=True))
|