"""Quick manual probe: python chat.py "سؤالك هنا" [model_dir]"""
import sys, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

question = sys.argv[1] if len(sys.argv) > 1 else "إذا كان لديك 1500 ريال وأنفقت 20% منها على الكتب، فكم تبقى معك؟"
model_dir = sys.argv[2] if len(sys.argv) > 2 else "./Nawah-Reasoning-v1"

tok = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.bfloat16).cuda().eval()

prompt = tok.apply_chat_template([{"role": "user", "content": question}], tokenize=False, add_generation_prompt=True)
ids = tok(prompt, return_tensors="pt").to("cuda")
out = model.generate(**ids, max_new_tokens=512, do_sample=False,
                     eos_token_id=[tok.eos_token_id, tok.convert_tokens_to_ids("<|im_end|>")],
                     pad_token_id=tok.pad_token_id)
print(tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=False).split("<|im_end|>")[0])