from transformers import AutoModelForCausalLM, AutoTokenizer
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--model_name", type=str, default="output_models/finetune_qwen_0.5B_shizhe_stage1_constant_lr_warmup_test")
args = parser.parse_args()

model_name = args.model_name

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 初始化对话历史
messages = [
    # {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}
]

print("聊天机器人已启动！输入 'exit' 退出对话。")
print("-" * 50)

while True:
    # 获取用户输入
    user_input = input("用户: ").strip()
    
    # 检查是否退出
    if user_input.lower() == "exit":
        print("再见！")
        break
    
    # 如果用户输入为空，继续下一轮
    if not user_input:
        continue
    
    # 添加用户消息到对话历史
    messages.append({"role": "user", "content": user_input})
    
    # 应用聊天模板
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    # 生成回复
    # generated_ids = model.generate(
    #     **model_inputs,
    #     max_new_tokens=512
    # )
    # generated_ids = [
    #     output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    # ]
    generated_ids = model.sample(
        model_inputs["input_ids"],
        n_steps=32,
        tokenizer=tokenizer,
    )
    
    response = tokenizer.batch_decode(generated_ids[:, model_inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
    print(f"AI: {response}")
    
    # 添加AI回复到对话历史
    messages.append({"role": "assistant", "content": response})
    
    print("-" * 50)
