数据准备阶段流程解析

代码仓库: rasbt/LLMs-from-scratch

本章节主要讲解了数据准备阶段的完整流程：
原始文本 → token → token ID → 向量
以下将逐步拆解每个阶段的核心内容。

1. 原始文本 → token 阶段

描述

在这一阶段，通过正则表达式将原始文本分割为 token（词或符号）。随后，使用 set 去重并通过 sorted() 排序，构建完整的词汇列表。

2. token → token ID 阶段

描述

通过 enumerate 结合字典推导式生成词表：

vocab = {token: integer for integer, token in enumerate(all_words)}

注意事项

手动构建词表时，通常需要加入特殊 token，例如：

  • <|unk|>：表示不在词表中的未知词。
  • <|endoftext|>：标记文本结束。
    这些特殊 token 用于避免模型在遇到词表外的 OOV（out-of-vocabulary）token 时出错。

而像 BPE（Byte Pair Encoding） 这样的分词算法则不需要 <|unk|>，因为它会将不认识的词拆分为更小的子词单位，确保所有内容都能映射到词表中。

2.1 BPE 简介

描述

BPE 是一种子词级别的分词算法。其核心思想是将不认识的 token 拆分为更小、更常见的子 token，从而全部映射到词表中。在项目中，我们通过以下代码使用 GPT-2 的 BPE 分词器：

tokenizer = tiktoken.get_encoding("gpt2")

这里的 gpt2 是基于 BPE 算法构建的分词器。

2.2 构建 Dataset 和 DataLoader

描述

  • Dataset：负责组织样本数据结构（token ID 的片段 + 目标）。
  • DataLoader：负责按批次从 Dataset 中提取数据（可控制 batch size、是否打乱等）。

示例代码：

from torch.utils.data import Dataset, DataLoader

class GPTDatasetV1(Dataset):
    def __init__(self, txt, tokenizer, max_length, stride):
        self.input_ids = []
        self.target_ids = []

        token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})

        for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i:i + max_length]
            target_chunk = token_ids[i + 1: i + max_length + 1]
            self.input_ids.append(torch.tensor(input_chunk))
            self.target_ids.append(torch.tensor(target_chunk))

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return self.input_ids[idx], self.target_ids[idx]

def create_dataloader_v1(txt, batch_size=4, max_length=256, 
                         stride=128, shuffle=True, drop_last=True,
                         num_workers=0):
    tokenizer = tiktoken.get_encoding("gpt2")
    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)

    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        num_workers=num_workers
    )

    return dataloader

2.2.1 滑动窗口

描述

滑动窗口是一种将长序列数据分割为多个小段的方式。例如，如果一段文本有 1000 个字符，而模型最多支持 200 个字符输入，可以设置：

max_length = 200
stride = 200

这样可以将文本切分为 5 个非重叠片段，每段作为一条训练数据。在此阶段，DataLoader 仅负责读取数据，尚未涉及将数据输入 LLM 进行推理。

若希望片段之间保留上下文信息，可设置 stride < max_length 来实现重叠窗口。

3. token ID → 向量阶段

Token ID to Vector Stage

3.1 Token 嵌入

描述

使用嵌入层将每个 token ID 映射为一个稠密向量：

embedding_layer = torch.nn.Embedding(vocab_size, output_dim)
max_length = 4
dataloader = create_dataloader_v1(
    raw_text, batch_size=8, max_length=max_length,
    stride=max_length, shuffle=False
)
data_iter = iter(dataloader)
inputs, targets = next(data_iter)

token_embeddings = embedding_layer(inputs)

3.2 位置嵌入

描述

如果序列中包含多个相同的 token，它们的 token embedding 会完全相同。但在自然语言中，同一词在不同位置的语义往往不同。为解决此问题，引入了位置嵌入：将每个 token 的 embedding 与对应位置的向量相加，得到最终的输入向量。最终的 input embedding 维度与 token embedding 一致。

4. 总结流程图

描述

原始文本  
  ↓ 分词 + 去重 + 编号  
Tokens → 词表 → Token IDs  
  ↓ 传入 embedding 层  
Token Embedding + Position Embedding  
  ↓  
最终输入向量