""" ============================================================================= 第 1 讲 · 代码 3：从零实现最简单的 BPE（Byte Pair Encoding） ============================================================================= 对应讲义：lectures/lecture01_overview_tokenization.md 第 10.1 ~ 10.5 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 上一个文件告诉我们： - 字节级分词器：永不遇到未知 token，但序列太长； - 词级分词器 ：序列短，但会遇到未知词。 BPE 想要两者兼得： 从"字节"出发（保证永远能编码任何文本）， 然后反复把**最常一起出现的相邻两个 token 合并成一个新 token**， 让常见的片段（如 "the "、"ing"）变成一个 token，罕见的片段仍然用字节表示。 这个文件是 BPE 的**教学版**：代码尽量短、尽量直白，和 CS336 课堂上的版本思路一致。 它有两个明显的缺点（我们会在演示中亲眼看到）： (1) 慢：每次合并都要扫描整个语料；编码时要遍历所有合并规则。 (2) 没有"预分词"：合并会跨越单词边界，学出 "e c" 这种奇怪的 token。 下一个文件 04_bpe_tokenizer.py 会逐一解决这两个问题。 ----------------------------------------------------------------------------- 【整体思路】 ----------------------------------------------------------------------------- 训练（train_bpe_naive）： 1. 把训练文本编码成 UTF-8 字节，得到初始 token 序列（每个 token 是 0~255）。 2. 重复 num_merges 次： a. 数一数所有"相邻 token 对"出现的次数（count_adjacent_pairs）； b. 找出出现次数最多的那一对 (a, b)； c. 给它分配一个新 id（256, 257, ...），把序列里所有 (a, b) 替换成新 id（merge）； d. 记录这条合并规则：(a, b) → 新 id；并在词表里记录新 id 对应的字节 = 字节(a) + 字节(b)。 编码（encode）： 把新文本变成字节，然后**按训练时的顺序**依次应用每条合并规则。 解码（decode）： 查词表把每个 id 变回字节，拼起来，再用 UTF-8 解码成字符串。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 03_bpe_from_scratch.py 只用到 Python 标准库。需要 data/sample_corpus.txt（本目录下已提供）。 """ # ============================================================================= # 第 0 部分：导入与工具函数 # ============================================================================= # 【Python 小知识】 # dataclass：一个"装饰器"，写在类定义上方（@dataclass）， # 能根据你列出的字段自动生成 __init__ 等方法，适合"只用来装数据"的类。 # time ：计时用。time.perf_counter() 返回一个高精度的"当前时刻"（秒）， # 两次调用相减就是经过的时间。 import time from dataclasses import dataclass from pathlib import Path CORPUS_PATH = Path(__file__).parent / "data" / "sample_corpus.txt" END_OF_TEXT = "<|endoftext|>" def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def load_corpus_documents(path=CORPUS_PATH): """读取语料并按 <|endoftext|> 切分成文档列表（详见 02 号文件中的同名函数）。""" path = Path(path) if not path.exists(): raise FileNotFoundError(f"找不到语料文件 {path}") docs = path.read_text(encoding="utf-8").split(END_OF_TEXT) return [d.strip() for d in docs if d.strip()] def is_chinese(doc): """ 粗略判断一篇文档是不是中文：只要含有一个常用汉字（U+4E00 ~ U+9FFF）就算。 【Python 小知识】 any(条件 for x in 序列)：只要有一个 x 让条件为真，就返回 True。 字符串之间可以用 <= 比较，比较的是码位大小。 """ return any("一" <= ch <= "鿿" for ch in doc) def split_train_test(docs): """ 把文档划分为训练集和测试集：留出最后一篇英文和最后一篇中文作为测试。 【思路】 评估分词器时，要用训练时**没见过**的文本，才能反映它在"新文本"上的真实压缩率。 （和机器学习里"不能在训练集上评估模型"是同一个道理。） 【返回】 (train_text, test_en, test_zh)，三个字符串。训练文档之间用换行连接。 """ english = [d for d in docs if not is_chinese(d)] chinese = [d for d in docs if is_chinese(d)] train_text = "\n".join(english[:-1] + chinese[:-1]) return train_text, english[-1], chinese[-1] def pretty_token(token_bytes): """ 把一个 token 的字节串变成适合打印的形式。 【思路】 BPE 的 token 是 bytes。大多数时候它能解码成正常文字（比如 b'the ' → 'the '）， 但有时它只是"半个汉字"（比如 b'\\xe4\\xbd'），无法解码。 能解码就显示文字，不能解码就显示原始字节。 【参数】 token_bytes : bytes 【返回】 str，例如 "'the '" 或 "b'\\xe4\\xbd'" 【Python 小知识】 repr(x) 返回 x 的"开发者视角"表示，字符串会带引号、特殊字符会被转义（如 '\\n'）， 这样空格和换行在打印时也能看清楚。 """ try: return repr(token_bytes.decode("utf-8")) except UnicodeDecodeError: return repr(token_bytes) # ============================================================================= # 第 1 部分：BPE 的两个基本操作 —— 数相邻对、合并 # ============================================================================= def count_adjacent_pairs(ids): """ 统计序列中每个"相邻 token 对"出现的次数。 【思路】 对序列 [t0, t1, t2, t3]，相邻对是 (t0,t1)、(t1,t2)、(t2,t3)。 用一个字典记录每一对出现了几次。 【参数】 ids : list[int]，token 序列。 【返回】 dict[tuple[int, int], int]，键是相邻对，值是出现次数。 字典的"插入顺序"就是每一对**第一次出现**的顺序（Python 3.7+ 的字典保证有序）。 【边界】 - 长度为 0 或 1 的序列没有相邻对，返回空字典 {}。 - 相邻对会**重叠计数**：[a, a, a] 中 (a, a) 计 2 次（位置 0-1 和 1-2）。 这和 merge 的行为（不重叠）不完全一致，但 BPE 的标准实现就是这样统计的。 【Python 小知识】 zip(ids, ids[1:]) 把"序列"和"去掉第一个元素的序列"并排对齐， 正好得到所有相邻对：zip([1,2,3], [2,3]) → (1,2), (2,3)。 counts.get(pair, 0) + 1：如果 pair 还没出现过，从 0 开始加 1。 """ counts = {} for pair in zip(ids, ids[1:]): counts[pair] = counts.get(pair, 0) + 1 return counts def merge(ids, pair, new_id): """ 把序列里所有出现的相邻对 pair 替换成 new_id。 【思路】 从左到右扫描：如果当前位置和下一个位置正好组成 pair， 就输出 new_id 并跳过 2 个位置；否则原样输出当前 token，前进 1 个位置。 【参数】 ids : list[int]，原序列（不会被修改）。 pair : tuple[int, int]，要合并的相邻对 (a, b)。 new_id : int，合并后的新 token id。 【返回】 list[int]，新序列。长度 = 原长度 - 被合并的次数。 【边界】 - 空序列 → 返回空列表。 - pair 不存在 → 返回与原序列相同内容的新列表。 - 重叠情况：merge([a, a, a], (a, a), X) → [X, a]，而不是 [X, X]。 从左到右"贪心"地合并，已经被合并用掉的 token 不能再参与下一次匹配。 【Python 小知识】 while 循环配合手动控制下标 i，比 for 循环更灵活：可以一次跳 1 步或 2 步。 i + 1 < len(ids) 的检查保证不会越界访问 ids[i + 1]。 """ new_ids = [] i = 0 while i < len(ids): if i + 1 < len(ids) and ids[i] == pair[0] and ids[i + 1] == pair[1]: new_ids.append(new_id) i += 2 else: new_ids.append(ids[i]) i += 1 return new_ids # ============================================================================= # 第 2 部分：训练 # ============================================================================= @dataclass class BPETokenizerParams: """ BPE 分词器的全部"参数"——训练完之后，分词器就由这两样东西完全决定。 【字段】 vocab : dict[int, bytes] 词表：token id → 这个 token 代表的字节串。 例如 {0: b'\\x00', ..., 116: b't', ..., 256: b'th', 257: b'the'} merges : dict[tuple[int, int], int] 合并规则：(a, b) → 新 id。 **字典的顺序就是训练时学到的顺序**，编码时必须按这个顺序应用。 """ vocab: dict merges: dict def choose_best_pair(counts, vocab): """ 从相邻对计数中选出"最好"的一对来合并。 【思路】 首先比较出现次数，次数最多的获胜。 如果有多个对出现次数相同（平局），需要一个**确定的**规则来打破平局， 否则同样的数据可能训练出不同的分词器。 这里采用 CS336 作业 1 的约定：选择**字典序更大**的那一对， 比较的是两个 token 的字节串 (vocab[a], vocab[b])。 （CS336 课堂代码用的是 max(counts, key=counts.get)，平局时取"最先出现"的那对。 对"the cat in the hat"这个例子，两种规则给出的结果恰好一样。） 【参数】 counts : dict[tuple[int,int], int]，count_adjacent_pairs 的输出，不能为空。 vocab : dict[int, bytes]，当前词表，用来查每个 token 的字节串。 【返回】 tuple[int, int]，被选中的相邻对。 【Python 小知识】 max(可迭代对象, key=函数)：对每个元素算出 key，返回 key 最大的那个元素。 这里 key 返回一个元组 (次数, 左边字节, 右边字节)。 元组的比较规则是"先比第一个元素，相同再比第二个……"， 所以相当于：先比次数，次数一样再比字节串的字典序。 bytes 之间的字典序：逐字节比较数值，例如 b't' (116) > b'h' (104)。 """ return max(counts, key=lambda p: (counts[p], vocab[p[0]], vocab[p[1]])) def train_bpe_naive(text, num_merges, verbose=False): """ 在一段文本上训练 BPE（教学版，没有预分词，没有特殊 token）。 【参数】 text : str，训练文本。 num_merges : int，要学习多少条合并规则（>= 0）。 最终词表大小 = 256 + 实际合并次数。 verbose : bool，True 时打印每一步的详细过程（适合很短的文本）。 【返回】 BPETokenizerParams 【边界】 - num_merges = 0：不合并，得到纯字节分词器。 - 文本太短、合并到后来已经没有相邻对了（整个文本变成了 1 个 token）： 提前停止，实际合并次数会少于 num_merges。 - 空文本：直接返回只有 256 个字节 token 的词表。 【复杂度】 每次合并都要扫描整个序列（数对 + 替换），所以总时间 ∝ num_merges × 文本长度。 对真实语料（几十 GB、几万次合并）这完全不可行 —— 04 号文件会优化。 """ if num_merges < 0: raise ValueError("num_merges 不能为负数") ids = list(text.encode("utf-8")) # 第 1 步：文本 → 字节序列 vocab = {i: bytes([i]) for i in range(256)} # 初始词表：256 个单字节 token merges = {} for step in range(num_merges): counts = count_adjacent_pairs(ids) # 第 2a 步：数相邻对 if not counts: if verbose: print(f" 第 {step + 1} 步：已经没有相邻对可以合并，提前结束。") break pair = choose_best_pair(counts, vocab) # 第 2b 步：选出最常见的一对 new_id = 256 + step # 第 2c 步：分配新 id ids = merge(ids, pair, new_id) # 替换 merges[pair] = new_id # 第 2d 步：记录规则 vocab[new_id] = vocab[pair[0]] + vocab[pair[1]] if verbose: # 把计数从大到小排序，只显示前 6 名 top = sorted(counts.items(), key=lambda kv: -kv[1])[:6] top_str = ", ".join(f"({pretty_token(vocab[a])},{pretty_token(vocab[b])})×{c}" for (a, b), c in top) print(f"\n 第 {step + 1} 次合并") print(f" 相邻对计数（前 6）：{top_str}") print(f" 选中：({pretty_token(vocab[pair[0]])}, {pretty_token(vocab[pair[1]])}) " f"出现 {counts[pair]} 次 → 新 token {new_id} = {pretty_token(vocab[new_id])}") print(f" 合并后序列（{len(ids)} 个 token）：{[vocab[i].decode('utf-8', 'replace') for i in ids]}") return BPETokenizerParams(vocab=vocab, merges=merges) # ============================================================================= # 第 3 部分：编码与解码 # ============================================================================= class NaiveBPETokenizer: """ 教学版 BPE 分词器：给定训练好的参数，实现 encode / decode。 【边界】 - 训练时没见过的字符（比如训练文本全是英文，却要编码中文）： 不会出错！它们会保持为单字节 token。这正是字节级 BPE 的优势——没有 。 - decode 遇到不在词表中的 id：抛出 ValueError。 """ def __init__(self, params): """ 【参数】 params : BPETokenizerParams（通常来自 train_bpe_naive） """ self.params = params def encode(self, text): """ 【思路】 1. 文本 → 字节序列； 2. 按训练时学到的顺序，依次应用每一条合并规则。 【复杂度】 对每条规则都扫描一遍序列：时间 ∝ 合并规则数 × 序列长度。 很多规则根本不会在这段文本中出现，却也被扫描了一遍 —— 这是浪费。 （CS336 课上也提到：作业里要"只遍历真正有关的合并规则"。） """ ids = list(text.encode("utf-8")) # 【Python 小知识】dict.items() 按插入顺序给出 (键, 值)，也就是训练时的合并顺序 for pair, new_id in self.params.merges.items(): ids = merge(ids, pair, new_id) return ids def decode(self, ids): """ 【思路】 先把每个 id 查表变成字节串，全部拼接后，**最后统一**做 UTF-8 解码。 （不能每个 token 单独解码：一个 token 可能只是半个汉字，见 01 号文件演示 5。） """ try: data = b"".join(self.params.vocab[i] for i in ids) except KeyError as e: raise ValueError(f"token id {e} 不在词表中") from e return data.decode("utf-8", errors="replace") def tokens_as_strings(self, text): """辅助函数：返回 encode 结果中每个 token 的可读形式，便于观察。""" return [self.params.vocab[i].decode("utf-8", errors="replace") for i in self.encode(text)] # ============================================================================= # 第 4 部分：演示 # ============================================================================= def demo_basic_operations(): """演示 1：count_adjacent_pairs 和 merge 的基本行为与边界。""" print_header("演示 1：两个基本操作 —— 数相邻对、合并") ids = [1, 2, 3, 1, 2] print(f" count_adjacent_pairs({ids}) = {count_adjacent_pairs(ids)}") print(f" merge({ids}, (1, 2), 99) = {merge(ids, (1, 2), 99)}") print("\n 【边界】") print(f" count_adjacent_pairs([]) = {count_adjacent_pairs([])}") print(f" count_adjacent_pairs([7]) = {count_adjacent_pairs([7])}") print(f" merge([], (1, 2), 99) = {merge([], (1, 2), 99)}") print(f" 重叠计数：count_adjacent_pairs([5, 5, 5]) = {count_adjacent_pairs([5, 5, 5])}") print(f" 重叠合并：merge([5, 5, 5], (5, 5), 99) = {merge([5, 5, 5], (5, 5), 99)}" " ← 从左到右贪心，不是 [99, 99]") # 自我检查 assert merge([1, 2, 3, 1, 2], (1, 2), 99) == [99, 3, 99] assert merge([5, 5, 5], (5, 5), 99) == [99, 5] assert merge([5, 5, 5, 5], (5, 5), 99) == [99, 99] print("\n ✓ 自我检查通过") def demo_train_trace(): """演示 2：在 "the cat in the hat" 上手动跟踪 3 次合并（CS336 课上的例子）。""" print_header('演示 2：在 "the cat in the hat" 上训练 3 次合并') text = "the cat in the hat" print(f" 训练文本：{text!r}") print(f" 初始字节序列（{len(text.encode())} 个 token）：{list(text.encode())}") params = train_bpe_naive(text, num_merges=3, verbose=True) print("\n 学到的合并规则（按顺序）：") for (a, b), new_id in params.merges.items(): print(f" ({a:>3}, {b:>3}) → {new_id} 即 {pretty_token(params.vocab[a])} + " f"{pretty_token(params.vocab[b])} = {pretty_token(params.vocab[new_id])}") return params def demo_encode_decode(params): """演示 3：用训练好的小分词器编码新文本。""" print_header("演示 3：编码与解码") tok = NaiveBPETokenizer(params) for text in ["the quick brown fox", "the hat", "他在看the cat"]: ids = tok.encode(text) print(f" 文本 {text!r}") print(f" ids = {ids}") print(f" tokens = {tok.tokens_as_strings(text)}") print(f" 字节数 {len(text.encode())} → token 数 {len(ids)}，" f"解码还原 = {tok.decode(ids) == text}") print("\n 注意第 3 个例子：训练时从没见过中文，但编码照样成功（中文退化为单字节 token），") print(" 并且解码能完美还原。字节级 BPE 永远不需要 。") # 边界：非法 id try: tok.decode([99999]) except ValueError as e: print(f"\n 【边界】decode([99999]) → ValueError：{e}") def demo_merge_order_matters(): """演示 4：为什么编码时必须按训练顺序应用合并规则。""" print_header("演示 4：合并顺序很重要") vocab = {i: bytes([i]) for i in range(256)} vocab[256] = b"ab" vocab[257] = b"bc" a, b, c = ord("a"), ord("b"), ord("c") right_order = BPETokenizerParams(vocab=vocab, merges={(a, b): 256, (b, c): 257}) wrong_order = BPETokenizerParams(vocab=vocab, merges={(b, c): 257, (a, b): 256}) print(" 合并规则：先 (a, b)→256，再 (b, c)→257") print(f" 按正确顺序编码 'abc' → {NaiveBPETokenizer(right_order).tokens_as_strings('abc')}") print(f" 按错误顺序编码 'abc' → {NaiveBPETokenizer(wrong_order).tokens_as_strings('abc')}") print(""" 两种切法都能正确解码回 'abc'，但它们是**不同的 token 序列**。 模型在训练时只见过按正确顺序切出来的序列，如果推理时换一种切法， 模型看到的就是它不熟悉的输入，效果会变差。 所以：分词器的"合并顺序"和词表一样，是分词器不可分割的一部分。""") def demo_bigger_corpus(): """演示 5：在真实一点的语料上训练，观察学到的 token，以及"跨词合并"的问题。""" print_header("演示 5：在小语料上训练 300 次合并") docs = load_corpus_documents() train_text, test_en, test_zh = split_train_test(docs) # 留出 1 篇英文 + 1 篇中文做测试 t0 = time.perf_counter() params = train_bpe_naive(train_text, num_merges=300) elapsed = time.perf_counter() - t0 tok = NaiveBPETokenizer(params) print(f" 训练文本 {len(train_text.encode()):,} 字节，300 次合并用时 {elapsed:.2f} 秒") first = [pretty_token(params.vocab[i]) for i in range(256, 256 + 30)] print(f"\n 最先学到的 30 个 token：\n {', '.join(first)}") longest = sorted(range(256, 256 + len(params.merges)), key=lambda i: -len(params.vocab[i]))[:12] print(f"\n 最长的 12 个 token：\n {', '.join(pretty_token(params.vocab[i]) for i in longest)}") # 找出"跨越单词边界"的 token：token 中间（不是开头）含有空格，或者字母后面跟着标点再跟字母等 cross = [params.vocab[i] for i in range(256, 256 + len(params.merges)) if b" " in params.vocab[i].strip(b" ") or params.vocab[i].endswith(b" ")] print(f"\n 其中有 {len(cross)} 个 token 跨越了单词边界（中间或末尾含空格），例如：") print(f" {', '.join(pretty_token(t) for t in cross[:12])}") print(" → 这些 token（如 'e '、'the '）把一个词的结尾和空格粘在一起。GPT-2 发现这会浪费词表，") print(" 例如 'dog'、'dog.'、'dog!' 会各占一个 token。解决办法是'预分词'（见 04 号文件）。") # 压缩率 print() for name, text in [("训练集", train_text), ("测试集·英文", test_en), ("测试集·中文", test_zh)]: ids = tok.encode(text) print(f" {name:<8}：{len(text.encode()):>6,} 字节 → {len(ids):>6,} token，" f"压缩率 {len(text.encode()) / len(ids):.2f}") assert tok.decode(ids) == text, "BPE 必须无损" print(" ✓ 所有文本都能无损还原") print(" → 英文测试集的压缩率和训练集相当（这些小故事用词高度重复）；") print(" 中文测试集明显更低：训练语料里中文只占一小部分，学到的中文 token 很少。") print(" 分词器会'偏向'它的训练数据 —— 讲义第 12 节'多语言不公平'就是这么来的。") def demo_speed(): """演示 6：教学版的训练和编码有多慢？""" print_header("演示 6：教学版 BPE 的速度") docs = load_corpus_documents() text = "\n".join(docs) num_bytes = len(text.encode()) print(f" 语料大小：{num_bytes:,} 字节") for n in [50, 100, 200, 400]: t0 = time.perf_counter() params = train_bpe_naive(text, num_merges=n) t_train = time.perf_counter() - t0 t0 = time.perf_counter() NaiveBPETokenizer(params).encode(text) t_encode = time.perf_counter() - t0 print(f" 合并 {n:>3} 次：训练 {t_train:5.2f} 秒，编码整个语料 {t_encode:5.2f} 秒") # 粗略外推：训练时间 ≈ (每次合并、每个字节的耗时) × 合并次数 × 语料字节数 sec_per_merge_per_byte = t_train / (n * num_bytes) gpt2_merges, owt_bytes = 50_000, 40e9 # GPT-2 约 5 万次合并；OpenWebText 约 40 GB est_years = sec_per_merge_per_byte * gpt2_merges * owt_bytes / (3600 * 24 * 365) print(f""" 训练时间和编码时间都大致与合并次数成正比。 粗略外推：每次合并每字节约 {sec_per_merge_per_byte:.1e} 秒； 若要像 GPT-2 那样在约 40 GB 文本上做 50,000 次合并，教学版大约需要 {est_years:.0f} 年。 （这只是数量级估计，而且还没算内存：40 GB 文本变成 Python 整数列表， 光是列表里每个元素 8 字节的指针就要约 320 GB 内存。） 04 号文件会把它优化到实用水平。""") def main(): demo_basic_operations() params = demo_train_trace() demo_encode_decode(params) demo_merge_order_matters() demo_bigger_corpus() demo_speed() print("\n全部演示运行完毕。下一个文件 04_bpe_tokenizer.py 是'作业级'的完整实现。") if __name__ == "__main__": main()