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))