AI 应用实例¶
本页收录可复用的 AI 应用脚本与踩坑:本地大模型流式推理 · 典型 bug 分析。命令优先、解释从简。
🟢 Qwen2.5 本地流式推理¶
transformers · Qwen2.5-Coder-7B · TextStreamer · CUDA 自动选择 · 逐 token 流式输出
本示例使用 transformers 加载 Qwen2.5-Coder-7B,实现一个支持逐字流式输出的本地对话程序,并在有 GPU 时自动使用 CUDA。
flowchart LR
A[用户输入 prompt] --> B[tokenizer 编码]
B --> C{检测设备}
C -->|有 CUDA| D[模型 → GPU]
C -->|无 CUDA| E[模型 → CPU]
D --> F[model.generate 单次生成]
E --> F
F --> G[TextStreamer 逐 token 打印]
G --> H[截掉输入 / 返回新文本]
依赖安装¶
运行前准备清单:
- 安装
torch与transformers - 下载 Qwen2.5-Coder-7B 到本地
- 修改脚本中的
model_path - 运行
python qwen_chat.py
完整脚本¶
qwen_chat.py
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import torch
import time
import sys
def print_slowly(text, delay=0.05):
"""逐字打印文本"""
for char in text:
sys.stdout.write(char)
sys.stdout.flush()
time.sleep(delay)
sys.stdout.write('\n')
def initialize_model(model_path):
print_slowly("正在加载模型,请稍候...", delay=0.1)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
device = "cuda" if torch.cuda.is_available() else "cpu" # (1)!
model = model.to(device)
print_slowly("模型加载完成!", delay=0.1)
return model, tokenizer, device
def generate_response(prompt, model, tokenizer, device):
# 构造输入
inputs = tokenizer(
prompt,
return_tensors="pt",
padding=True,
return_attention_mask=True
)
# 将输入移到同样的设备上
inputs = {k: v.to(device) for k, v in inputs.items()}
# 使用 TextStreamer 实现逐 token 流式输出,无需手动循环
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) # (2)!
# 单次生成,由 streamer 负责实时打印
with torch.no_grad():
outputs = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
max_new_tokens=1024,
do_sample=True,
temperature=0.7,
pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id, # (3)!
streamer=streamer,
)
# 去掉输入部分,只保留新生成的文本
new_tokens = outputs[0][inputs["input_ids"].shape[1]:] # (4)!
return tokenizer.decode(new_tokens, skip_special_tokens=True)
def main():
model_path = "/home/liujingbo/Qwen2.5-Coder-7B"
model, tokenizer, device = initialize_model(model_path)
print_slowly("\n欢迎使用AI助手!输入 'quit' 或 'exit' 结束对话", delay=0.05)
while True:
# 获取用户输入
user_input = input("\n你: ")
# 检查是否退出
if user_input.lower() in ['quit', 'exit']:
print_slowly("再见!", delay=0.001)
break
# 如果输入为空,继续下一轮
if not user_input.strip():
continue
# 生成回答(由 TextStreamer 实时打印)
print("\nAI: ", end='')
generate_response(user_input, model, tokenizer, device)
print()
if __name__ == "__main__":
main()
torch.cuda.is_available()自动判断有无 GPU,选择cuda/cpu。TextStreamer在generate过程中逐 token 解码打印,无需生成结束后再手动模拟。- 部分模型未设置 pad token,回退到
eos_token_id避免报错。 - 截掉输入部分,只保留新生成的内容。
关键说明¶
| 要点 | 说明 |
|---|---|
| 流式输出 | 由 TextStreamer 完成,它在 model.generate 过程中逐 token 解码并打印,无需在生成结束后再逐字模拟 |
pad_token_id |
使用 tokenizer.pad_token_id or tokenizer.eos_token_id,避免部分模型未设置 pad token 时报错 |
| 截断输入 | 通过 new_tokens = outputs[0][inputs["input_ids"].shape[1]:] 截掉输入部分,只返回新生成内容 |
🐛 原脚本存在的问题¶
model.generate · for 循环陷阱 · do_sample · max_new_tokens
原始脚本在生成回答时使用了如下写法,存在严重 bug:
原脚本(有 bug)
generated_text = ''
with torch.no_grad():
for _ in range(200):
outputs = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
max_new_tokens=10000,
do_sample=True,
temperature=0.7,
pad_token_id=tokenizer.pad_token_id
)
outputs = outputs.cpu()
new_token = tokenizer.decode(outputs[0][-1], skip_special_tokens=True)
print(new_token, end="", flush=True)
generated_text += new_token
问题根因
model.generate每次都是基于同一份输入完整生成一遍文本,并不存在「接着上一次继续生成」的语义。 外层for _ in range(200)只是把整段生成重复 200 次。- 每次生成的随机性(
do_sample=True)不同,且只取outputs[0][-1](最后一个 token)打印, 导致输出是 200 个互不相关的片段拼接,逻辑上等同于死循环式重复生成。 max_new_tokens=10000又让每次生成的代价极高,进一步放大问题。
修正方式是单次调用 model.generate 并传入 TextStreamer,由框架负责真正的流式输出,代码见上文。