# %% Cell 10: Quick Test Inference
messages = [
    {"role": "user", "content": "What does Charaka say about the qualities of a good physician in Sutrasthana?"}
]

inputs = tokenizer.apply_chat_template(
    messages, tokenize=True, return_tensors="pt", add_generation_prompt=True
).to("cuda")

print("Generating...")
outputs = model.generate(
    inputs,
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"\n{'='*60}")
print("RESPONSE")
print(f"{'='*60}")
print(response)