# 修改前（去网上找）：
# def __init__(self, model_name="Qwen/Qwen3-0.6B"):

# 修改后（去你电脑的本地找）：
class QwenChatbot:
    def __init__(self, model_name=r"D:\models\Qwen3-0.6B"):  # ⬅️ 这里换成你电脑上的真实路径
        print("🤖 机器人正在加载大脑和翻译官，请稍等...")
        # ... 后面的代码保持不变

# 第一步：从工具箱里拿出我们需要的工具
# AutoModelForCausalLM：这是机器人的“大脑”
# AutoTokenizer：这是机器人的“翻译官”，负责把我们的话变成机器能懂的数字
from modelscope import AutoModelForCausalLM, AutoTokenizer
import torch  # 引入 torch 库，这是用来控制机器人运行速度的“加速器”

# 第二步：定义一个机器人的“图纸”（在编程里叫“类”）
class QwenChatbot:
    # 这是机器人的“出生说明书”（初始化方法）
    # 当你第一次启动它时，它就会按照这里的步骤准备自己
    def __init__(self, model_name="Qwen/Qwen3-0.6B"):
        print("🤖 机器人正在加载大脑和翻译官，请稍等...")
        
        # 1. 加载翻译官
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # 2. 加载大脑（⚠️ 重点优化：加了 device_map="auto" 和 torch_dtype）
        # device_map="auto"：意思是“请自动帮我找一台合适的电脑（或显卡）来运行我”
        # torch_dtype=torch.bfloat16：意思是“请用省内存的模式加载大脑”，这样小电脑也能跑得动！
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, 
            device_map="auto", 
            torch_dtype=torch.bfloat16
        )
        
        # 3. 让大脑进入“考试模式”（⚠️ 重点优化：加了 eval()）
        # eval() 的意思是告诉大脑：“现在是回答问题时间，不要做多余的动作（比如随机忘记东西）”
        self.model.eval() 
        
        # 4. 准备一个“小本本”（历史记录），用来记住我们之前聊过什么
        self.history = []
        print("✅ 机器人准备就绪！可以开始聊天啦！\n")

    # 这是机器人的“说话功能”
    def generate_response(self, user_input):
        # 1. 把新说的话和“小本本”上的旧话拼在一起
        messages = self.history + [{"role": "user", "content": user_input}]

        # 2. 让翻译官把拼好的话，变成机器人能看懂的格式
        text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        # 3. 把文本变成数字（⚠️ 重点优化：把数字放到和大脑相同的电脑设备上）
        # 如果不加 .to(self.model.device)，大脑和数字不在同一个房间，机器人会报错罢工哦！
        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
        
        # 4. 大脑开始思考并回答问题（最多想 32768 个词）
        response_ids = self.model.generate(
            **inputs, 
            max_new_tokens=32768
        )[0][len(inputs.input_ids[0]):].tolist()
        
        # 5. 把机器人想出来的数字，翻译回我们能看懂的文字
        response = self.tokenizer.decode(response_ids, skip_special_tokens=True)

        # 6. 把刚才的对话记到“小本本”上（⚠️ 重点优化：防止小本本太厚被撑爆）
        # 如果聊得太久，小本本（历史记录）超过了 10 条，我们就把最旧的一条擦掉！
        self.history.append({"role": "user", "content": user_input})
        self.history.append({"role": "assistant", "content": response})
        
        # 如果历史记录超过 20 个（也就是 10 轮对话），就删掉最早的 2 个
        if len(self.history) > 20:
            self.history = self.history[2:] 

        # 7. 把机器人的回答交给你
        return response

# 第三步：启动机器人！（只有直接运行这个文件时，下面的代码才会执行）
if __name__ == "__main__":
    # 1. 制造一个机器人
    chatbot = QwenChatbot()

    # 2. 问它第一个问题（默认开启深度思考模式）
    user_input_1 = "How many r's in strawberries?"
    print(f"👦 你: {user_input_1}")
    response_1 = chatbot.generate_response(user_input_1)
    print(f"🤖 机器人: {response_1}")
    print("-" * 30)

    # 3. 问它第二个问题，并加上 /no_think 标签
    # 意思是：“这个问题很简单，你不用深度思考，直接回答我！”
    user_input_2 = "Then, how many r's in blueberries? /no_think"
    print(f"👦 你: {user_input_2}")
    response_2 = chatbot.generate_response(user_input_2)
    print(f"🤖 机器人: {response_2}") 
    print("-" * 30)

    # 4. 问它第三个问题，并加上 /think 标签
    # 意思是：“这个问题很难，请你开启深度思考模式，仔细想一想再回答！”
    user_input_3 = "Really? /think"
    print(f"👦 你: {user_input_3}")
    response_3 = chatbot.generate_response(user_input_3)
    print(f"🤖 机器人: {response_3}")