推理模型从零构建 第2章：使用Qwen3生成文本

from pathlib import Path
from reasoning_from_scratch.qwen3 import Qwen3Tokenizer, download_qwen3_small

# 下载 tokenizer 文件
download_qwen3_small(kind="base", tokenizer_only=True, out_dir="qwen3")

# 加载 tokenizer
tokenizer_path = Path("qwen3") / "tokenizer-base.json"
tokenizer = Qwen3Tokenizer(tokenizer_file_path=tokenizer_path)

# encode: 文本 → token ID 列表
prompt = "Explain large language models."
input_token_ids = tokenizer.encode(prompt)

for i in input_token_ids:
    print(f"{i} --> {tokenizer.decode([i])}")
# 840  --> Ex
# 20772 --> plain
# 3460  -->  large
# 4128  -->  language
# 4119  -->  models
# 13    --> .

# decode: token ID 列表 → 文本（无损还原）
print(tokenizer.decode(input_token_ids))
# "Explain large language models."

import torch

def get_device(enable_tensor_cores=True):
    if torch.cuda.is_available():
        device = torch.device("cuda")
        # 启用 TF32 加速矩阵乘法（Ampere 及以上架构）
        if enable_tensor_cores:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    elif torch.xpu.is_available():
        device = torch.device("xpu")
    else:
        device = torch.device("cpu")
    return device

device = get_device()  # 自动选择最佳设备

from reasoning_from_scratch.qwen3 import (
    Qwen3Model, QWEN_CONFIG_06_B, download_qwen3_small
)

# 下载预训练权重（~1.5 GB）
download_qwen3_small(kind="base", tokenizer_only=False, out_dir="qwen3")

# 实例化模型并加载权重
model_path = Path("qwen3") / "qwen3-0.6B-base.pth"
model = Qwen3Model(QWEN_CONFIG_06_B)
model.load_state_dict(torch.load(model_path))
model.to(device)

@torch.inference_mode()
def generate_text_basic_stream(model, token_ids, max_new_tokens, eos_token_id=None):
    model.eval()

    for _ in range(max_new_tokens):
        # 前向传播，取最后一个位置的 logits
        out = model(token_ids)[:, -1]
        # 贪心解码：选概率最高的 token
        next_token = torch.argmax(out, dim=-1, keepdim=True)

        # 遇到终止符则停止
        if eos_token_id is not None and torch.all(next_token == eos_token_id):
            break

        yield next_token  # 流式输出每个 token
        # 将新 token 拼接到输入序列末尾
        token_ids = torch.cat([token_ids, next_token], dim=1)


# 使用示例
prompt = "Explain large language models in a single sentence."
input_ids = torch.tensor(tokenizer.encode(prompt), device=device).unsqueeze(0)

for token in generate_text_basic_stream(model, input_ids, max_new_tokens=100,
                                         eos_token_id=tokenizer.eos_token_id):
    print(tokenizer.decode(token.squeeze(0).tolist()), end="", flush=True)
# → Large language models are artificial intelligence systems that can
#   understand, generate, and process human language, enabling them to
#   perform a wide range of tasks, from answering questions to writing
#   articles, and even creating creative content.

from reasoning_from_scratch.qwen3 import KVCache

@torch.inference_mode()
def generate_text_basic_stream_cache(model, token_ids, max_new_tokens,
                                      eos_token_id=None):
    model.eval()
    cache = KVCache(n_layers=model.cfg["n_layers"])  # 初始化缓存
    model.reset_kv_cache()

    # Prefill：一次性处理整个 prompt，填充 cache
    out = model(token_ids, cache=cache)[:, -1]

    for _ in range(max_new_tokens):
        next_token = torch.argmax(out, dim=-1, keepdim=True)

        if eos_token_id is not None and torch.all(next_token == eos_token_id):
            break

        yield next_token
        # 只传入新生成的 1 个 token，cache 中已有历史 K/V
        out = model(next_token, cache=cache)[:, -1]

# 一行代码开启编译优化
model_compiled = torch.compile(model)

# 首次运行会触发编译（较慢），后续运行大幅加速
# 使用方式和原模型完全一样
for token in generate_text_basic_stream_cache(
    model=model_compiled,
    token_ids=input_ids,
    max_new_tokens=100,
    eos_token_id=tokenizer.eos_token_id
):
    print(tokenizer.decode(token.squeeze(0).tolist()), end="", flush=True)