from unsloth import FastLanguageModel, get_chat_template
import torch
# 1. 加载合并后的模型和分词器
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "/home/merged_model1",
    load_in_4bit = False,  # 如果保存时是4bit合并，这里要True；16bit则可False或不写
)
tokenizer = get_chat_template(tokenizer, chat_template="chatml", map_eos_token=True)

# 2. 构造输入
messages = [
    {"role": "user", "content": "你好，假设你是一个五年级的数学老师，你的学生很调皮而且不集中注意力听课，你要怎么讲课才能让他们学会鸡兔同笼的解题方法？"}
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text=prompt, return_tensors="pt").to(model.device)#注意这里要注明是text，要不然模型会认为你应该输入图片，然后报错

# 3. 推理生成
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=512)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(result)