本地部署qwen3 1.7B 聊天机器人全流程代码
·
本地部署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)
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐



所有评论(0)