第 4 章 数据处理流水线

第 4 章 数据处理流水线

模型吃的是张量，不是文本。本章讲解原始文本如何经过「下载 → 分词 → 切窗」变成训练样本，以及本项目的两个特色：样本索引返回 与 训练/验证集切分。

4.1 数据流水线总览

整条流水线的目标：把人类可读的文本，变成形状为 [batch, seq_len] 的整数张量，供模型训练。

4.2 数据来源与三级回退策略

本项目使用 Karpathy 的经典 tiny_shakespeare 数据集（约 1MB 的莎士比亚文本），它小到单机几分钟就能训完，又足够展示语言模型的核心行为。

获取数据时采用三级回退策略，保证在各种网络环境下都能跑起来：

核心代码逻辑：

from pathlib import Path
import urllib.request

_TINY_SHAKESPEARE_URL = (
    "https://raw.githubusercontent.com/karpathy/char-rnn/"
    "master/data/tinyshakespeare/input.txt"
)

def get_dataset(cache_path="data/tiny_shakespeare.txt") -> str:
    cache_path = Path(cache_path)
    cache_path.parent.mkdir(parents=True, exist_ok=True)

    # 1) 命中本地缓存，直接读
    if cache_path.exists():
        return cache_path.read_text(encoding="utf-8")

    text = None
    # 2) 尝试 HuggingFace datasets
    try:
        from datasets import load_dataset
        ds = load_dataset("tiny_shakespeare", split="train")
        text = "\n".join(ds["text"])
    except Exception:
        pass

    # 3) 回退：GitHub 直链
    if text is None:
        with urllib.request.urlopen(_TINY_SHAKESPEARE_URL) as resp:
            text = resp.read().decode("utf-8")

    cache_path.write_text(text, encoding="utf-8")
    return text
​

设计要点：try/except 包住可选依赖（datasets），缺失时优雅降级；最后兜底用标准库 urllib 直接下载。这样即便在一个只装了 torch 和 tiktoken 的极简环境，也能拿到数据。

4.3 分词：tiktoken 与 p50k_base

文本不能直接喂给模型，要先切成 token（子词）并映射为整数 id。本项目用 tiktoken 的 p50k_base 编码——这正是 OpenAI GPT-2 使用的分词器。

import tiktoken

enc = tiktoken.get_encoding("p50k_base")

ids = enc.encode("To be, or not to be")
print(ids)            # [1438, 307, 11, ...]
print(enc.decode(ids))  # 还原回 "To be, or not to be"
​

几个特点：

• 词表大小固定 50257：与配置中的 vocab_size 对齐。
• BPE 算法：高频词整体编码，低频词拆成子词，兼顾词表大小与覆盖率。
• Rust 实现：比纯 Python 分词器快一个数量级。

全局复用编码器：tiktoken 内部有较大的 BPE 表，反复创建开销不小。本项目用一个全局变量复用同一个编码器实例：

_ENCODER = tiktoken.get_encoding("p50k_base")

def get_encoder():
    return _ENCODER
​

4.4 自回归样本构造：x 与 y

语言模型学的是「预测下一个 token」。给定一段文本，我们把它切成「输入 x」和「目标 y」——y 就是 x 左移一位。

文本片段（block_size+1 个 token）:  [t0, t1, t2, t3, t4]
                                       └──── x ────┘ └ y 位置
x = [t0, t1, t2, t3]      ← 模型输入
y = [t1, t2, t3, t4]      ← 每个位置的目标是「下一个 token」
​

这样在 x 的每个位置 i，模型都要预测 y[i]，等于一次前向就能在 block_size 个位置上同时学习。

数据集类的实现

import torch
from torch.utils.data import Dataset

class TextDataset(Dataset):
    def __init__(self, tokens: list, block_size: int, name="train"):
        self.block_size = block_size
        self.tokens = tokens
        # 每个样本需要 block_size+1 个 token
        self.n_samples = max(0, len(tokens) - block_size - 1)

    def __len__(self):
        return self.n_samples

    def __getitem__(self, idx):
        chunk = self.tokens[idx : idx + self.block_size + 1]
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y
​

注意 chunk[:-1] 与 chunk[1:] 的对称关系——这正是「左移一位」的实现。

4.5 特色之一：样本索引返回

普通语言模型数据集只返回 (x, y)。但本项目的数据集额外返回样本索引 idx：

def __getitem__(self, idx):
    chunk = self.tokens[idx : idx + self.block_size + 1]
    x = torch.tensor(chunk[:-1], dtype=torch.long)
    y = torch.tensor(chunk[1:], dtype=torch.long)
    return x, y, idx   # ← 多了一个 idx
​

为什么要返回索引？这是为教师输出缓存服务的。当启用缓存时，教师会预先把每个样本的输出 logits 存到磁盘，训练时按 idx 查表取出，不必每次都跑教师前向。这个机制在《第 10 章》详细讲，这里只需记住：数据集返回 idx，让缓存成为可能。

如果不启用缓存，这个 idx 会被训练循环忽略，不影响普通流程。

4.6 DataLoader：批处理与加速

Dataset 一次只给一个样本，训练时需要把它们打包成批次，并做乱序、预取等优化。这是 DataLoader 的工作：

from torch.utils.data import DataLoader

def build_dataloader(dataset, batch_size, num_workers=0, shuffle=True):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        drop_last=True,                              # 丢弃不完整批次
        pin_memory=torch.cuda.is_available(),         # GPU 加速
    )
​

几个关键参数：

Windows 注意：若 num_workers > 0，必须保证训练入口在 if __name__ == "__main__": 中（本项目已遵守），否则会触发多进程递归启动问题。

4.7 特色之二：训练集与验证集切分

评估模型需要一份「模型没见过的」验证集。本项目从全文末尾切出约 10% 作为验证集，避免与训练集重叠：

为什么从末尾切，而不是随机切？因为文本是连续的，随机切会让训练集和验证集混入彼此的上下文，导致验证集「泄题」。从末尾切一段干净的，能更真实反映泛化能力。

实现：

def build_train_eval_loaders(text, block_size, batch_size, num_workers=0, eval_ratio=0.1):
    split_char = int(len(text) * (1.0 - eval_ratio))
    train_text = text[:split_char]
    eval_text = text[split_char:]

    train_dataset = TextDataset.from_text(train_text, block_size, name="train")
    eval_dataset = TextDataset.from_text(eval_text, block_size, name="eval")

    train_loader = build_dataloader(train_dataset, batch_size, num_workers, shuffle=True)
    eval_loader = build_dataloader(eval_dataset, batch_size, num_workers, shuffle=False)
    return train_loader, eval_loader, train_dataset
​

验证集的用途在《第 8 章 评估与对比》中会展开——我们会在上面计算困惑度、Top-1 准确率，并对比师生表现。

4.8 端到端验证：看一眼真实数据

把上面所有环节串起来，打印一个真实批次：

from dataset import get_dataset, build_train_eval_loaders

text = get_dataset()
train_loader, eval_loader, _ = build_train_eval_loaders(
    text, block_size=32, batch_size=2
)

# 取一个批次
for x, y, idx in train_loader:
    print("x 形状:", x.shape, "dtype:", x.dtype)   # [2, 32] long
    print("y 形状:", y.shape)
    print("idx:", idx.tolist())
    print("x 第一个样本前 8 个 token:", x[0, :8].tolist())
    print("y 第一个样本前 8 个 token:", y[0, :8].tolist())
    # 验证 y 是 x 左移一位
    assert torch.equal(x[0, 1:8], y[0, :7])
    print("✓ y 确实是 x 左移一位")
    break
​

预期输出：

x 形状: torch.Size([2, 32]) dtype: torch.int64
y 形状: torch.Size([2, 32])
idx: [1234, 567]
x 第一个样本前 8 个 token: [464, 1234, ...]
y 第一个样本前 8 个 token: [1234, ...]
✓ y 确实是 x 左移一位
​

本章小结

• 数据流水线：文本 → 分词 → token 序列 → 切窗 → 样本 → 批次。
• 三级回退策略保证数据在各种网络环境都能获取。
• tiktoken 的 p50k_base 分词，词表 50257，与 GPT-2 对齐。
• 自回归样本：y 是 x 左移一位，每个位置学「预测下一个 token」。
• 数据集额外返回 idx，为教师输出缓存埋下伏笔。
• 从文本末尾切出验证集，避免与训练集重叠。

动手实验：把 block_size 改成 16，重新打印一个批次，观察样本数量和形状的变化，理解切窗逻辑。

下一站：数据准备好了，主角登场。在《第 5 章 教师与学生模型》中，我们看教师如何加载、学生如何初始化，以及为什么教师要被冻结。