跳转至

AI 应用实例

本页收录可复用的 AI 应用脚本与踩坑:本地大模型流式推理 · 典型 bug 分析。命令优先、解释从简。

  • 本地流式推理


    transformers + TextStreamer 实现 Qwen2.5 逐 token 输出

    查看

  • 原脚本 Bug


    for 循环里反复调用 model.generate 的重复生成陷阱

    查看


🟢 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[截掉输入 / 返回新文本]

依赖安装

pip install torch transformers
uv add torch transformers

运行前准备清单:

  • 安装 torchtransformers
  • 下载 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()
  1. torch.cuda.is_available() 自动判断有无 GPU,选择 cuda / cpu
  2. TextStreamergenerate 过程中逐 token 解码打印,无需生成结束后再手动模拟。
  3. 部分模型未设置 pad token,回退到 eos_token_id 避免报错。
  4. 截掉输入部分,只保留新生成的内容。

关键说明

要点 说明
流式输出 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,由框架负责真正的流式输出,代码见上文。