import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

model = AutoModelForCausalLM.from_pretrained("thanghf/math_model_v1",torch_dtype=torch.bfloat16,device_map='auto')
tokenizer = AutoTokenizer.from_pretrained("thanghf/math_model_v1")
model.eval()
streamer = TextStreamer(tokenizer)
prompt = """Gieo hai con súc xắc cân đối và đồng chất. Xác suất để tổng số chấm trên mặt xuất hiện của hai con súc xắc bằng 7 là:"""
messages = [
    {"role": "system", "content": "Please reason step by step, and put your final answer within \\boxed{}."},
    {"role": "user", "content": prompt}
]

text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(device)

generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=4096,
    streamer=streamer
)