本地部署qwen3 1.7B 聊天机器人全流程代码 

要求:魔法

import os
import time  # 导入时间模块用于计时
from transformers import AutoModelForCausalLM, AutoTokenizer

class QwenChatbot:
    def __init__(self, model_name="Qwen/Qwen3-1.7B"):
        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        # 加载模型
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype="auto",
            device_map="auto"
        )
        # 初始化对话历史列表
        self.history = []
        
    def generate_response(self, user_input):
        # 固定启用思考
        enable_thinking = True
        
        # 将新输入添加到消息列表(历史对话+新输入)
        messages = self.history + [{"role": "user", "content": user_input}]
        
        # 应用聊天模板
        text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=enable_thinking  # 始终为True
        )
        
        # 聊天模板token化
        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
        
        # 记录生成开始时间
        start_time = time.time()
        
        # 生成回复
        generated_ids = self.model.generate(
            **model_inputs,
            max_new_tokens=32768  # 控制最大生成token数
        )
        
        # 记录生成结束时间
        end_time = time.time()
        # 计算耗时(保留2位小数)
        elapsed_time = round(end_time - start_time, 2)
        
        # 提取模型新生成的部分(排除原始输入)
        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
        
        # 分离思考内容和最终回复(Qwen3的特殊功能)
        try:
            # 查找思考结束标记的位置(151668是Qwen3中思考结束符号的token ID)
            index = len(output_ids) - output_ids[::-1].index(151668)
        except ValueError:
            # 如果没有找到思考标记,全部作为最终回复
            index = 0
            
        # 解码思考内容和最终回复
        thinking_content = self.tokenizer.decode(
            output_ids[:index], 
            skip_special_tokens=True
        ).strip("\n")
        
        content = self.tokenizer.decode(
            output_ids[index:], 
            skip_special_tokens=True
        ).strip("\n")
        
        # 更新对话历史,保存本轮对话
        self.history.append({"role": "user", "content": user_input})
        self.history.append({"role": "assistant", "content": content})
        
        # 返回思考内容、回复内容和耗时
        return thinking_content, content, elapsed_time

# 示例用法
if __name__ == "__main__":
    # 创建聊天机器人实例
    chatbot = QwenChatbot()
    
    print("Qwen3聊天机器人已启动,输入'退出'结束对话...")
    
    while True:
        # 获取用户输入
        user_input = input("\n你: ")
        
        # 检查是否退出
        if user_input.lower() in ["退出", "q", "quit"]:
            print("聊天结束,再见!")
            break
        
        # 生成回复(包含耗时记录)
        thinking, response,耗时 = chatbot.generate_response(user_input)
        
        # 显示结果(包含耗时)
        print(f"\n本轮回复耗时: {耗时} 秒")
        print("\n机器人思考过程:")
        print(thinking)
        print("\n机器人回复:")
        print(response)

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐