import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. 加载训练好的模型和分词器
model_path = "./model"  # 对应你训练保存的路径

print("正在加载模型...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()  # 切换到评估模式
print("模型加载完成！")


# 2. 定义推理对话函数
def generate_response(user_message, max_new_tokens=128):
    prompt = f"Human: {user_message}\nBot:"

    # 编码输入
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    prompt_len = inputs["input_ids"].shape[1]  # 获取输入的 token 数量

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

    # 【关键修改】直接切片去掉输入 Prompt 部分的 token，只解码模型新生成的内容
    generated_tokens = outputs[0][prompt_len:]
    response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()

    # 如果模型生成的开头依然带有 "Bot:" 或 "ot:"，用正则或 strip 清理掉
    import re
    response = re.sub(r'^(Bot:|bot:|ot:|\s)+', '', response).strip()

    return response

# 3. 简单的交互式对话循环
if __name__ == "__main__":
    print("\n--- GPT-2 中文聊天机器人已启动 (输入 'exit' 或 'quit' 退出) ---")
    while True:
        user_input = input("\nHuman: ")
        if user_input.strip().lower() in ["exit", "quit"]:
            print("再见！")
            break

        if not user_input.strip():
            continue

        bot_reply = generate_response(user_input).replace(" ", "")
        print(f"Bot: {bot_reply}")