""" ============================================================================= 第 1 讲 · 代码 4：完整的字节级 BPE 分词器（作业 1 风格） ============================================================================= 对应讲义：lectures/lecture01_overview_tokenization.md 第 10.6 ~ 10.9 节、第 15 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 03 号文件的教学版 BPE 能跑，但离"能用"还差三件事： 问题 1：合并会跨越单词边界，学出 'e '、'. They ' 这种 token，浪费词表。 → 解决：预分词（pre-tokenization）。先用 GPT-2 的正则表达式把文本切成 "类似单词"的小块，BPE 只在每个小块**内部**合并。 问题 2：没有特殊 token。训练语料里的 <|endoftext|> 会被当成普通字符拆开、参与合并。 → 解决：先按特殊 token 把文本切开；特殊 token 整体作为一个 id，不参与合并。 问题 3：太慢。训练时每次合并都扫描整个语料；编码时要遍历所有合并规则。 → 解决： 训练：(a) 相同的"词"只处理一次（乘以出现次数）； (b) 每次合并后只更新**受影响**的词，而不是全部重算。 编码：(a) 每个预分词块只看它**内部存在**的相邻对，按优先级合并； (b) 缓存：同一个词编码过一次就记住结果。 这正是 CS336 作业 1 第 2 部分要你实现的东西。 【如果你正在做作业】建议先自己写，卡住了再来看这里的思路。直接抄代码学不到东西。 ----------------------------------------------------------------------------- 【整体结构】 ----------------------------------------------------------------------------- 第 1 部分：预分词（GPT-2 正则表达式） 第 2 部分：特殊 token 的切分 第 3 部分：训练 —— 高效版 train_bpe / train_bpe_from_text，以及用于对拍的慢速参考版 第 4 部分：BPETokenizer 类 —— encode / decode / encode_iterable / save / load 第 5 部分：演示与自我检查 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 04_bpe_tokenizer.py 只需要 Python 标准库即可运行。 可选：安装 regex 库（pip install regex）后，将使用与 GPT-2 **完全相同**的预分词正则； 没装的话，会用一个用标准库 re 写的近似版本（对英文和中文结果相同，见第 1 部分说明）。 """ # ============================================================================= # 第 0 部分：导入 # ============================================================================= # 【Python 小知识】 # json ：读写 JSON 格式（一种通用的文本数据格式），用来保存/加载分词器。 # tempfile ：创建临时文件夹，演示保存/加载时用，程序结束后自动删除。 # defaultdict ：一种字典，访问不存在的键时会自动创建默认值。 # defaultdict(int)[新键] → 0；defaultdict(set)[新键] → set()（空集合）。 # 这样就不用每次都写 "if key not in d: d[key] = 0"。 import json import re import tempfile import time from collections import Counter, defaultdict from pathlib import Path # 【Python 小知识】try / except ImportError：尝试导入一个可选的库，没装也不报错。 try: import regex # 第三方库，支持 \p{L} 这类 Unicode 属性写法 HAS_REGEX = True except ImportError: regex = None HAS_REGEX = False 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 pretty_token(token_bytes): """token 字节串的可读形式：能解码就显示文字，否则显示字节（见 03 号文件）。""" try: return repr(token_bytes.decode("utf-8")) except UnicodeDecodeError: return repr(token_bytes) # ============================================================================= # 第 1 部分：预分词（pre-tokenization） # ============================================================================= # 【思路】 # 在 BPE 合并之前，先把文本切成"类似单词"的小块（pre-token）。之后： # - 合并只发生在每个小块内部，不会跨越小块边界； # - 相同的小块（比如 " the" 出现了 1 万次）只需要处理一次，再乘以次数。 # # GPT-2 使用的正则表达式（CS336 作业 1 也要求用它）： # # '(?:[sdmt]|ll|ve|re) 英文缩写：'s 'd 'm 't 'll 've 're 单独成块 # （只认小写和 ASCII 撇号，所以 "DON'T" 会被切成 DON / ' / T） # | ?\p{L}+ 可选的一个空格 + 一串字母（\p{L} = 任何语言的字母，包括汉字） # | ?\p{N}+ 可选的一个空格 + 一串数字（GPT-2 不限长度；GPT-4 改成最多 3 位） # | ?[^\s\p{L}\p{N}]+ 可选的一个空格 + 一串"既不是空白、也不是字母数字"的符号（标点、emoji） # |\s+(?!\S) 一串空白，但如果后面紧跟着非空白字符，就留下最后一个空白字符 # （(?!\S) 叫"负向前瞻"：要求后面不是非空白字符，本身不消耗字符） # |\s+ 剩下的空白 # # 关键设计：**空格跟着后面的词走**。" world" 是一个块（空格在前）， # 所以 "hello world" → ["hello", " world"]。这就是为什么 GPT 的词表里有大量以空格开头的 token。 # # "a b"（3 个空格）→ ["a", " ", " b"]： # \s+ 本可以吃掉 3 个空格，但 (?!\S) 要求后面不能紧跟非空白字符， # 于是回退到只吃 2 个空格，把最后 1 个空格留给 " b"。 # 这样不论前面有多少空格，"b" 总是以 " b" 的形式出现，词表不会因为空格数量不同而膨胀。 # # 【边界：没有安装 regex 库时】 # Python 标准库 re 不支持 \p{L}、\p{N}。我们用下面的等价写法近似： # \p{L} ≈ [^\W\d_] （"单词字符"去掉数字和下划线 = 字母） # \p{N} ≈ \d （十进制数字） # [^\s\p{L}\p{N}] ≈ [^\s\w] 或 _ # 对英文、中文、常见标点和 emoji，两者切分结果完全相同； # 只在极少数字符上有差别（例如上标数字 '²'，GPT-2 视为数字，近似版视为字母）。 # **两者都保证切分是无损的**：所有块拼起来一定等于原文（每个字符都会被某个分支匹配到）。 GPT2_PATTERN = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" FALLBACK_PATTERN = r"""'(?:[sdmt]|ll|ve|re)| ?[^\W\d_]+| ?\d+| ?(?:[^\s\w]|_)+|\s+(?!\S)|\s+""" if HAS_REGEX: PRETOKENIZE_RE = regex.compile(GPT2_PATTERN) PRETOKENIZER_NAME = "regex 库 + GPT-2 原版正则" else: PRETOKENIZE_RE = re.compile(FALLBACK_PATTERN) PRETOKENIZER_NAME = "标准库 re + 近似正则（pip install regex 可切换为原版）" def pretokenize(text): """ 把一段（不含特殊 token 的）文本切成预分词块。 【参数】 text : str。可以是空字符串（返回空列表）。 【返回】 list[str]，例如 "Hello, world!" → ['Hello', ',', ' world', '!'] 【边界】 - 保证无损：''.join(pretokenize(text)) == text 对任何输入都成立。 - 这个函数**不认识**特殊 token。如果直接把 "<|endoftext|>" 传进来， 它会被切成 ['<|', 'endoftext', '|>']。所以调用前要先用 split_on_special_tokens 切开。 【Python 小知识】 pattern.findall(text) 返回所有不重叠的匹配子串。 正则里的 (?:...) 是"非捕获分组"，只用来把几个选项括起来，不影响 findall 的返回值。 """ return PRETOKENIZE_RE.findall(text) # ============================================================================= # 第 2 部分：特殊 token # ============================================================================= # 【思路】 # 特殊 token（如 <|endoftext|>）是人为规定的"控制符号"，用来标记文档边界等信息。 # 它们必须满足： # (1) 永远作为一个整体，对应唯一的一个 id，不能被拆开； # (2) 不参与 BPE 合并统计 —— 否则会学出 "<|endoftext|>Once" 这种跨文档的 token； # (3) 训练时，特殊 token 两侧的文本要分开处理，合并不能跨越它。 # 实现方法：先用特殊 token 把文本"切开"，特殊 token 和普通文本分别处理。 def split_on_special_tokens(text, special_tokens): """ 按特殊 token 把文本切成若干段，并标记每一段是不是特殊 token。 【参数】 text : str special_tokens : list[str] 或 tuple[str]，可以为空。 【返回】 list[tuple[str, bool]]：[(片段, 是否特殊 token), ...]，空片段会被丢弃。 例如 split_on_special_tokens("Hi<|endoftext|>Yo", ["<|endoftext|>"]) → [("Hi", False), ("<|endoftext|>", True), ("Yo", False)] 【边界】 - 没有特殊 token：整段作为普通文本返回（空文本返回空列表）。 - **特殊 token 互相重叠**：例如同时有 "<|endoftext|>" 和 "<|endoftext|><|endoftext|>"。 正则的"或"（|）会选择**第一个**能匹配的选项，所以必须把**长的放在前面**， 否则长的那个永远不会被匹配到。我们按长度从大到小排序来保证这一点。 - 特殊 token 里可能含有正则的元字符（比如 | 和 .），必须用 re.escape 转义， 否则 "<|endoftext|>" 里的 | 会被当成"或"。 【Python 小知识】 re.split(pattern, text)：按匹配到的位置切开字符串。 如果 pattern 外面包了一层**捕获括号** (...)，匹配到的分隔符本身也会保留在结果里， 并且位于结果列表的奇数下标位置： re.split("(,)", "a,b") → ['a', ',', 'b'] """ if not special_tokens: return [(text, False)] if text else [] ordered = sorted(set(special_tokens), key=len, reverse=True) # 长的在前 pattern = "(" + "|".join(re.escape(s) for s in ordered) + ")" parts = re.split(pattern, text) result = [] for i, part in enumerate(parts): if part == "": continue is_special = (i % 2 == 1) # 奇数下标 = 被捕获的分隔符 = 特殊 token result.append((part, is_special)) return result # ============================================================================= # 第 3 部分：训练 # ============================================================================= def merge_word(word, pair, new_token): """ 在一个"词"（token 列表）中，把所有相邻的 pair 合并成 new_token。 与 03 号文件的 merge 完全相同，只是这里的 token 是 bytes 而不是 int。 用 bytes 表示 token 的好处：比较大小（打破平局）时可以直接比字节串。 【参数】 word : list[bytes]，例如 [b't', b'h', b'e'] pair : tuple[bytes, bytes]，例如 (b't', b'h') new_token : bytes，通常就是 pair[0] + pair[1]，例如 b'th' 【返回】 list[bytes]，例如 [b'th', b'e'] 【边界】与 03 号文件相同：从左到右、不重叠地合并。 """ out = [] i = 0 while i < len(word): if i + 1 < len(word) and word[i] == pair[0] and word[i + 1] == pair[1]: out.append(new_token) i += 2 else: out.append(word[i]) i += 1 return out def count_pretokens(text, special_tokens): """ 统计每个预分词块（以 UTF-8 字节形式）出现的次数。 【思路】 这是训练加速的第一步：真实语料里 " the" 可能出现上亿次， 但我们只需要把它作为"一个词、出现上亿次"来处理，而不是处理上亿遍。 对于英文，不同的预分词块数量通常比总块数小好几个数量级。 【参数】 text : str，整个训练语料 special_tokens : list[str] 【返回】 Counter，键是 bytes（预分词块的 UTF-8 编码），值是出现次数。 特殊 token 本身不会被计入。 """ counts = Counter() for piece, is_special in split_on_special_tokens(text, special_tokens): if is_special: continue # 特殊 token 不参与合并统计 for pt in pretokenize(piece): counts[pt.encode("utf-8")] += 1 return counts def _validate_training_args(vocab_size, special_tokens): """ 检查训练参数是否合法，返回去重后的特殊 token 列表。（函数名以下划线开头，表示"内部使用"。） 【边界】 - vocab_size 必须 >= 256 + 特殊 token 个数，否则连基础字节都放不下。 - 特殊 token 重复出现时只保留一次（保持原顺序）。 - 空字符串不能作为特殊 token（它会在每个位置都"匹配"）。 """ # 【Python 小知识】dict.fromkeys(列表) 会按顺序去重（字典的键不能重复，且保持插入顺序） specials = list(dict.fromkeys(special_tokens)) if any(s == "" for s in specials): raise ValueError("特殊 token 不能是空字符串") min_size = 256 + len(specials) if vocab_size < min_size: raise ValueError(f"vocab_size={vocab_size} 太小：至少需要 256（字节）+ {len(specials)}（特殊 token）= {min_size}") return specials def train_bpe_from_text(text, vocab_size, special_tokens=(), verbose=False): """ 在一段文本上训练字节级 BPE（高效版）。 【参数】 text : str，训练语料。 vocab_size : int，最终词表大小（包括 256 个字节、特殊 token、合并产生的 token）。 special_tokens : list[str]，特殊 token，例如 ["<|endoftext|>"]。 verbose : bool，True 时每 100 次合并打印一次进度。 【返回】 (vocab, merges) vocab : dict[int, bytes]，token id → 字节串 id 分配方式：0~255 是单字节；接着是特殊 token；最后是合并产生的新 token。 merges : list[tuple[bytes, bytes]]，按学习顺序排列的合并规则。 这与 CS336 作业 1 要求的 train_bpe 返回格式一致。 【算法】 1. 预分词并计数：{" the": 1523, " cat": 87, ...} 2. 把每个不同的块拆成单字节列表：" the" → [b' ', b't', b'h', b'e'] 3. 建立两张表： pair_counts : 相邻对 → 在整个语料中出现的总次数（= Σ 所在词的出现次数） pair_to_words : 相邻对 → 包含它的词的编号集合（"倒排索引"） 4. 重复直到词表满： a. 选出次数最多的相邻对（平局时选字典序更大的，作业 1 的约定）； b. 只对包含它的那些词（查倒排索引）执行合并，并**增量更新**两张表： 先减去这个词旧的所有相邻对，再加上合并后新的所有相邻对。 【为什么增量更新是对的】 一个词的合并只会影响这个词自己的相邻对，不会影响其他词。 所以其他词的计数完全不用动，只需要修正被合并的词。 【边界】 - 语料太小、所有词都已经合并成单个 token、没有相邻对可合并：提前停止， 此时词表会比 vocab_size 小（会打印提示）。 - 选最大值这一步每次都扫描所有相邻对，复杂度 O(不同相邻对的个数)。 更快的做法是用"堆"（heapq）配合"懒删除"，但代码会复杂很多，这里没有采用。 """ specials = _validate_training_args(vocab_size, special_tokens) # ---- 初始化词表 ---- vocab = {i: bytes([i]) for i in range(256)} for s in specials: vocab[len(vocab)] = s.encode("utf-8") num_merges = vocab_size - len(vocab) # ---- 第 1、2 步：预分词、计数、拆成单字节 ---- pretoken_counts = count_pretokens(text, specials) # 【Python 小知识】bytes 可以直接遍历，得到的是整数；bytes([b]) 把整数变回长度为 1 的 bytes words = [[bytes([b]) for b in pt] for pt in pretoken_counts] # 每个不同的块 freqs = list(pretoken_counts.values()) # 对应的出现次数 # ---- 第 3 步：建立相邻对计数和倒排索引 ---- pair_counts = defaultdict(int) pair_to_words = defaultdict(set) for wi, (word, freq) in enumerate(zip(words, freqs)): for pair in zip(word, word[1:]): pair_counts[pair] += freq pair_to_words[pair].add(wi) # ---- 第 4 步：循环合并 ---- merges = [] for step in range(num_merges): if not pair_counts: print(f" [提示] 第 {step} 次合并时已无相邻对可合并，提前结束（词表大小 {len(vocab)} < {vocab_size}）") break # 4a. 选出最好的一对。key 是 (次数, 相邻对)；相邻对是 (bytes, bytes)， # 元组比较 = 先比左边字节串、再比右边字节串 → "字典序更大"的规则。 best = max(pair_counts, key=lambda p: (pair_counts[p], p)) new_token = best[0] + best[1] merges.append(best) vocab[len(vocab)] = new_token # 4b. 只更新包含 best 的词。 # 【Python 小知识】list(...) 先复制一份，因为循环中会修改 pair_to_words[best] 这个集合， # "一边遍历一边修改"同一个集合会报错。 for wi in list(pair_to_words[best]): old_word = words[wi] new_word = merge_word(old_word, best, new_token) freq = freqs[wi] # 减去旧词的所有相邻对 for pair in zip(old_word, old_word[1:]): pair_counts[pair] -= freq if pair_counts[pair] == 0: del pair_counts[pair] # 计数归零就删掉，避免 max 时扫描无用的项 pair_to_words[pair].discard(wi) # discard：元素不存在也不报错（remove 会报错） # 加上新词的所有相邻对 for pair in zip(new_word, new_word[1:]): pair_counts[pair] += freq pair_to_words[pair].add(wi) words[wi] = new_word pair_to_words.pop(best, None) # best 已经不存在于任何词中了，清理掉 if verbose and (step + 1) % 100 == 0: print(f" 已完成 {step + 1}/{num_merges} 次合并，最新 token：{pretty_token(new_token)}") return vocab, merges def train_bpe(input_path, vocab_size, special_tokens=()): """ 从文件训练 BPE。与 CS336 作业 1 要求的函数签名一致。 【参数】 input_path : str 或 Path，UTF-8 编码的文本文件。 vocab_size : int special_tokens : list[str] 【返回】 (vocab, merges)，同 train_bpe_from_text。 【边界 / 大文件怎么办】 这里一次性把整个文件读进内存。对于几 GB 的真实语料（例如 OpenWebText）， 作业里推荐的做法是： 1. 把文件切成若干块，**切分点必须落在特殊 token 上**（这样不会把一个词切成两半）； 2. 用多进程（multiprocessing）并行地对每一块做预分词和计数； 3. 把各块的 Counter 加起来，再做合并循环。 预分词通常是训练中最耗时的部分，并行化能带来接近线性的加速。 """ text = Path(input_path).read_text(encoding="utf-8") return train_bpe_from_text(text, vocab_size, special_tokens) def train_bpe_reference(text, vocab_size, special_tokens=()): """ 【慢速参考实现】用于"对拍"（验证高效版的正确性）。 【思路】 每一轮都从头重新统计所有相邻对，然后合并所有词。 逻辑最简单、最不容易写错，但每次合并都要处理全部数据。 如果高效版和它的结果完全一致，我们就有信心高效版的增量更新没有写错。 （这是写算法时非常实用的技巧：先写一个慢但显然正确的版本，再用它检验快的版本。） """ specials = _validate_training_args(vocab_size, special_tokens) vocab = {i: bytes([i]) for i in range(256)} for s in specials: vocab[len(vocab)] = s.encode("utf-8") pretoken_counts = count_pretokens(text, specials) words = {pt: [bytes([b]) for b in pt] for pt in pretoken_counts} merges = [] while len(vocab) < vocab_size: pair_counts = Counter() for pt, word in words.items(): for pair in zip(word, word[1:]): pair_counts[pair] += pretoken_counts[pt] if not pair_counts: break best = max(pair_counts, key=lambda p: (pair_counts[p], p)) merges.append(best) vocab[len(vocab)] = best[0] + best[1] words = {pt: merge_word(w, best, best[0] + best[1]) for pt, w in words.items()} return vocab, merges # ============================================================================= # 第 4 部分：BPETokenizer 类 # ============================================================================= class BPETokenizer: """ 字节级 BPE 分词器：encode / decode / encode_iterable / save / load。 【用法】 vocab, merges = train_bpe("data/sample_corpus.txt", 500, ["<|endoftext|>"]) tok = BPETokenizer(vocab, merges, special_tokens=["<|endoftext|>"]) ids = tok.encode("Once upon a time<|endoftext|>") text = tok.decode(ids) 【编码算法（为什么"按优先级合并"等价于"按顺序应用所有规则"）】 对每个预分词块： 1. 拆成单字节列表； 2. 在块内所有相邻对中，找出**合并优先级最高**（训练时最早学到，rank 最小）的那一对； 3. 把块内所有这一对都合并； 4. 重复 2~3，直到块内没有任何相邻对在合并规则表里。 这与 03 号文件"按训练顺序逐条应用所有规则"的结果完全相同，但快得多： 只看块内实际存在的几个相邻对，而不是遍历几万条规则。 等价性的理由：合并 (A,B)→C 只会让 C 和它的左右邻居形成新的相邻对， 而包含 C 的规则一定比 (A,B) 学得更晚（C 是 (A,B) 之后才出现的）， 所以"当前块里 rank 最小的相邻对"一定就是"按顺序扫描时下一条会生效的规则"。 【边界】 - 训练时没见过的字符：退化为单字节 token，永远不会出错（没有 ）。 - 特殊 token：如果构造时传入了 special_tokens，文本中出现的特殊 token 会被编码为单个 id； 如果没传，它们会被当作普通文字拆开编码。 （安全提示：如果用户输入里混进了 "<|endoftext|>"，而你把它当特殊 token 处理， 用户就能"伪造"文档边界。tiktoken 默认会拒绝编码用户文本中的特殊 token，就是这个原因。） - 极少数情况下，两条不同的合并规则可能产生相同的字节串（比如 (a,bc) 和 (ab,c) 都得到 abc）。 此时词表里同一字节串有两个 id，编码时统一使用编号最小的那个。 """ def __init__(self, vocab, merges, special_tokens=None): """ 【参数】 vocab : dict[int, bytes] merges : list[tuple[bytes, bytes]]，按训练顺序排列 special_tokens : list[str] 或 None。如果某个特殊 token 不在 vocab 中，会自动追加到词表末尾。 """ self.vocab = dict(vocab) # 复制一份，避免修改调用者传进来的字典 self.merges = list(merges) self.special_tokens = list(dict.fromkeys(special_tokens or [])) # 字节串 → id 的反向表。按 id 从小到大插入，setdefault 保证重复字节串保留最小 id。 # 【Python 小知识】d.setdefault(k, v)：如果 k 不在 d 中，就设 d[k] = v；已存在则不改动。 self.bytes_to_id = {} for i in sorted(self.vocab): self.bytes_to_id.setdefault(self.vocab[i], i) for s in self.special_tokens: b = s.encode("utf-8") if b not in self.bytes_to_id: new_id = max(self.vocab) + 1 if self.vocab else 0 self.vocab[new_id] = b self.bytes_to_id[b] = new_id # 合并规则的优先级：越早学到，rank 越小，优先级越高 self.merge_ranks = {pair: rank for rank, pair in enumerate(self.merges)} # 缓存：预分词块（bytes）→ 编码结果（list[int]） # 自然语言里常见词反复出现，缓存能大幅提速。 # 【边界】缓存没有大小上限；在超大语料上长期运行可能占用较多内存。 self._cache = {} # ------------------------------------------------------------------ # 编码 # ------------------------------------------------------------------ def _encode_pretoken(self, pretoken_bytes): """ 对单个预分词块做 BPE 合并，返回 id 列表。（内部方法，外部请用 encode。） 【参数】 pretoken_bytes : bytes，一个预分词块的 UTF-8 编码，例如 b' the'。 【返回】 list[int] """ cached = self._cache.get(pretoken_bytes) if cached is not None: return cached parts = [bytes([b]) for b in pretoken_bytes] while len(parts) >= 2: # 找出块内 rank 最小（优先级最高）的相邻对 best_pair, best_rank = None, None for pair in zip(parts, parts[1:]): rank = self.merge_ranks.get(pair) if rank is not None and (best_rank is None or rank < best_rank): best_pair, best_rank = pair, rank if best_pair is None: break # 没有可以合并的相邻对了 parts = merge_word(parts, best_pair, best_pair[0] + best_pair[1]) try: ids = [self.bytes_to_id[p] for p in parts] except KeyError as e: raise ValueError(f"合并结果 {e} 不在词表中：vocab 和 merges 不匹配（文件损坏或来自不同的分词器？）") from e self._cache[pretoken_bytes] = ids return ids def encode(self, text): """ 把字符串编码为 token id 列表。 【参数】 text : str。可以为空（返回 []）。 【返回】 list[int] """ ids = [] for piece, is_special in split_on_special_tokens(text, self.special_tokens): if is_special: ids.append(self.bytes_to_id[piece.encode("utf-8")]) else: for pt in pretokenize(piece): ids.extend(self._encode_pretoken(pt.encode("utf-8"))) return ids def encode_iterable(self, iterable): """ 逐块编码一个"可迭代的字符串序列"（比如打开的文件对象，每次给一行），逐个产出 token id。 【思路】 对几 GB 的大文件，不能一次性读进内存再 encode。 这个方法每次只处理一小块，内存占用与文件大小无关。 【参数】 iterable : 任何能逐个给出 str 的东西，例如 open(path, encoding="utf-8")、字符串列表。 【返回】 生成器（generator），逐个产出 int。可以用 list(...) 收集，或在 for 循环中逐个使用。 【边界】 按块（行）切开后分别编码，**在极少数情况下**与"整体 encode"结果不同： 预分词的 \\s+(?!\\S) 规则需要"向后看一个字符"，而切块边界截断了这种前瞻。 例如整体文本 "x\\n\\n\\ny" 预分词为 ['x', '\\n\\n', '\\n', 'y']， 而按行切成 "x\\n"、"\\n"、"\\n"、"y" 后，得到 ['x', '\\n', '\\n', '\\n', 'y']。 如果词表里有 "\\n\\n" 这个 token，两种编码的 id 序列就会不同。 严格一致的做法是缓存一段"尾巴"，等看到下一块再决定如何切分；这里为简洁没有实现。 【Python 小知识】 函数里用了 yield 就变成"生成器"：每次 yield 一个值后暂停， 等调用方要下一个值时再继续运行。yield from 列表 = 逐个 yield 列表中的元素。 """ for chunk in iterable: yield from self.encode(chunk) # ------------------------------------------------------------------ # 解码 # ------------------------------------------------------------------ def decode(self, ids): """ 把 token id 列表解码为字符串。 【思路】先拼接所有字节，最后统一 UTF-8 解码；非法字节用 U+FFFD（�）替换。 【边界】id 不在词表中：抛出 ValueError。 """ try: data = b"".join(self.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 token_strings(self, text): """辅助：返回 encode 结果中每个 token 的可读形式（便于观察切分）。""" return [self.vocab[i].decode("utf-8", errors="replace") for i in self.encode(text)] # ------------------------------------------------------------------ # 保存 / 加载 # ------------------------------------------------------------------ def save(self, path): """ 把分词器保存为一个 JSON 文件。 【思路】 JSON 不能直接存 bytes，所以我们把每个字节串转成十六进制字符串（b'th' → "7468"）。 GPT-2 的做法不同：它用一张"字节 → 可打印 Unicode 字符"的映射表（bytes_to_unicode）， 把每个 token 存成一个看起来像文字的字符串（例如空格显示为 'Ġ'）。 你在 GPT-2 的词表文件里看到的 "Ġthe" 就是 " the"。 """ data = { "vocab": {str(i): b.hex() for i, b in self.vocab.items()}, "merges": [[a.hex(), b.hex()] for a, b in self.merges], "special_tokens": self.special_tokens, } Path(path).write_text(json.dumps(data, ensure_ascii=False, indent=1), encoding="utf-8") @classmethod def load(cls, path): """ 从 save 保存的 JSON 文件加载分词器。 【Python 小知识】 @classmethod 表示这是一个"类方法"：用 BPETokenizer.load(...) 调用， 第一个参数 cls 就是类本身，cls(...) 相当于 BPETokenizer(...)，用来创建新实例。 bytes.fromhex("7468") → b'th'。 """ data = json.loads(Path(path).read_text(encoding="utf-8")) vocab = {int(i): bytes.fromhex(h) for i, h in data["vocab"].items()} merges = [(bytes.fromhex(a), bytes.fromhex(b)) for a, b in data["merges"]] return cls(vocab, merges, special_tokens=data.get("special_tokens", [])) # ============================================================================= # 第 5 部分：演示与自我检查 # ============================================================================= def is_chinese(doc): """含有常用汉字就算中文文档。""" return any("一" <= ch <= "鿿" for ch in doc) def load_train_test(): """ 读取语料，留出最后一篇英文和最后一篇中文作为测试文本。 【返回】 (train_text, test_en, test_zh) train_text 中文档之间用 "<|endoftext|>" 连接（和原始文件格式一致）， 这样训练时也能演示特殊 token 的处理。 """ raw = CORPUS_PATH.read_text(encoding="utf-8") docs = [d.strip() for d in raw.split(END_OF_TEXT) if d.strip()] english = [d for d in docs if not is_chinese(d)] chinese = [d for d in docs if is_chinese(d)] train_docs = english[:-1] + chinese[:-1] train_text = END_OF_TEXT.join(train_docs) return train_text, english[-1], chinese[-1] def demo_pretokenize(): """演示 1：预分词。""" print_header("演示 1：预分词（GPT-2 正则）") print(f" 当前使用：{PRETOKENIZER_NAME}\n") examples = [ "Hello, world! I'm here.", "a b", "DON'T don't", "12345 apples cost $3.50", " return x", "你好，世界！", "hi 🌍!!", "x\n\n\ny", ] for s in examples: pts = pretokenize(s) assert "".join(pts) == s, "预分词必须无损" print(f" {s!r:28} → {pts}") print("\n ✓ 所有例子：''.join(预分词结果) == 原文（无损）") # 在整个语料上检查无损性 text = CORPUS_PATH.read_text(encoding="utf-8") assert "".join(pretokenize(text)) == text print(" ✓ 在整个语料文件上同样无损") # 如果装了 regex 库，比较原版正则和近似正则是否一致 if HAS_REGEX: fallback = re.compile(FALLBACK_PATTERN) same = fallback.findall(text) == pretokenize(text) print(f" 原版 GPT-2 正则 与 标准库近似正则 在本语料上的切分是否完全一致：{same}") tricky = "x² ½ Ⅻ" print(f" 少数特殊字符上的差异示例 {tricky!r}：原版 {pretokenize(tricky)}，近似 {fallback.findall(tricky)}") print("\n 【边界】pretokenize 不认识特殊 token：", pretokenize("end<|endoftext|>start")) def demo_special_tokens(): """演示 2：按特殊 token 切分，以及"长的优先"的重要性。""" print_header("演示 2：特殊 token 的切分") s = "Hello<|endoftext|>World" print(f" {s!r}\n → {split_on_special_tokens(s, [END_OF_TEXT])}") specials = [END_OF_TEXT, END_OF_TEXT * 2] s2 = "a<|endoftext|><|endoftext|>b<|endoftext|>c" print(f"\n 特殊 token 有重叠：{specials}") print(f" {s2!r}") print(f" 正确（长的优先）→ {split_on_special_tokens(s2, specials)}") # 故意用"短的优先"来演示错误 wrong_pattern = "(" + "|".join(re.escape(t) for t in sorted(specials, key=len)) + ")" wrong = [p for p in re.split(wrong_pattern, s2) if p] print(f" 错误（短的优先）→ {wrong} ← 双重 token 永远匹配不到") print(f"\n 【边界】没有特殊 token 时：{split_on_special_tokens('abc', [])}；空文本：{split_on_special_tokens('', [END_OF_TEXT])}") def demo_train(): """演示 3：训练，并与慢速参考实现对拍。返回训练好的分词器。""" print_header("演示 3：训练 BPE（vocab_size = 500）") train_text, _, _ = load_train_test() specials = [END_OF_TEXT] t0 = time.perf_counter() vocab, merges = train_bpe_from_text(train_text, vocab_size=500, special_tokens=specials) t_fast = time.perf_counter() - t0 print(f" 训练文本 {len(train_text.encode()):,} 字节；高效版用时 {t_fast:.3f} 秒，得到 {len(merges)} 条合并规则") print(f" 词表：{len(vocab)} = 256（字节）+ {len(specials)}（特殊 token）+ {len(merges)}（合并）") print(f" 特殊 token 的 id：{[i for i, b in vocab.items() if b == END_OF_TEXT.encode()]}") print("\n 最先学到的 20 条合并：") for k, (a, b) in enumerate(merges[:20]): print(f" {k + 1:>2}. {pretty_token(a):>8} + {pretty_token(b):<8} → {pretty_token(a + b)}") new_tokens = [a + b for a, b in merges] longest = sorted(new_tokens, key=len, reverse=True)[:12] print(f"\n 最长的 12 个 token：{', '.join(pretty_token(t) for t in longest)}") # 预分词之后，不应再出现"中间或末尾带空格"的 token（空格只可能出现在开头） bad = [t for t in new_tokens if b" " in t.lstrip(b" ") and t.strip(b" ")] print(f" 中间或末尾含空格的 token 个数：{len(bad)} ← 03 号教学版里有几十个，预分词后没有了") # 对拍：高效版 vs 慢速参考版 t0 = time.perf_counter() ref_vocab, ref_merges = train_bpe_reference(train_text, vocab_size=500, special_tokens=specials) t_ref = time.perf_counter() - t0 assert ref_merges == merges and ref_vocab == vocab, "高效版与参考版结果不一致！" print(f"\n ✓ 对拍通过：高效版与慢速参考版得到完全相同的 vocab 和 merges") print(f" 参考版用时 {t_ref:.3f} 秒，高效版快了约 {t_ref / t_fast:.0f} 倍") # 边界：vocab_size 太小 try: train_bpe_from_text("abc", vocab_size=100) except ValueError as e: print(f"\n 【边界】vocab_size=100 → ValueError：{e}") return BPETokenizer(vocab, merges, special_tokens=specials) def demo_encode_decode(tok): """演示 4：编码与解码，包括各种边界输入。""" print_header("演示 4：编码与解码") _, test_en, test_zh = load_train_test() examples = [ "Once upon a time, there was a little cat.", "the cat in the hat", "Hello, 🌍! 你好!", "小青蛙和小蜗牛一起看彩虹。", "def f(x):\n return x * 2", "The end.<|endoftext|>A new story begins.", "", " ", ] for s in examples: ids = tok.encode(s) back = tok.decode(ids) assert back == s, f"往返失败：{s!r}" shown = tok.token_strings(s) print(f" {s!r}") print(f" → {len(ids)} 个 token：{shown if len(shown) <= 16 else shown[:16] + ['...']}") print("\n ✓ 所有例子都能无损往返（包括空字符串、纯空格、emoji、特殊 token）") print("\n 测试集压缩率（字节 / token）：") for name, text in [("英文", test_en), ("中文", test_zh)]: ids = tok.encode(text) print(f" {name}：{len(text.encode()):,} 字节 → {len(ids):,} token，压缩率 {len(text.encode()) / len(ids):.2f}") # 特殊 token 的"安全"问题 plain = BPETokenizer(tok.vocab, tok.merges, special_tokens=None) s = "user says <|endoftext|>" print(f"\n 同一句 {s!r}：") print(f" 启用特殊 token → 末尾是 1 个 id：{tok.encode(s)[-1:]}") print(f" 不启用 → 被当成普通文字：{plain.token_strings(s)[-4:]}") # 边界：非法 id try: tok.decode([10**9]) except ValueError as e: print(f"\n 【边界】decode([10**9]) → ValueError：{e}") def demo_encode_iterable(tok): """演示 5：流式编码大文件。""" print_header("演示 5：encode_iterable —— 流式编码") # 【Python 小知识】with open(...) as f：打开文件，with 块结束时自动关闭文件。 # 对文件对象做 for 循环，每次得到一行（包含行尾的 '\n'）。 with open(CORPUS_PATH, encoding="utf-8") as f: streamed = list(tok.encode_iterable(f)) whole = tok.encode(CORPUS_PATH.read_text(encoding="utf-8")) print(f" 逐行流式编码：{len(streamed):,} 个 token；整体编码：{len(whole):,} 个 token；完全相同：{streamed == whole}") s = "x\n\n\ny" lines = s.splitlines(keepends=True) # ['x\n', '\n', '\n', 'y'] print(f"\n 【边界】{s!r} 整体预分词：{pretokenize(s)}") print(f" 按行 {lines} 分别预分词：{[p for line in lines for p in pretokenize(line)]}") print(" 预分词结果不同；只有当词表里有 '\\n\\n' 这类 token 时，最终 id 才会不同。") def demo_save_load(tok): """演示 6：保存与加载。""" print_header("演示 6：保存与加载") # 【Python 小知识】TemporaryDirectory 创建一个临时文件夹，with 块结束后自动删除，不会弄乱你的目录 with tempfile.TemporaryDirectory() as d: path = Path(d) / "tokenizer.json" tok.save(path) size_kb = path.stat().st_size / 1024 tok2 = BPETokenizer.load(path) text = "Once upon a time<|endoftext|>小明笑了。" assert tok2.encode(text) == tok.encode(text) and tok2.vocab == tok.vocab print(f" 保存为 JSON（约 {size_kb:.0f} KB），重新加载后编码结果完全相同 ✓") def measure_throughput(tok, text, use_cache=True, repeats=10): """ 测量编码吞吐量（字节/秒）。 【思路】 每一轮都**新建**一个分词器（缓存是空的），编码一遍文本，重复多轮取总时间。 这样缓存只能在"一遍文本内部"发挥作用（同一个词第二次出现时命中）， 不会因为"同一段文本编码了很多遍"而虚高。 【参数】 tok : BPETokenizer，只用它的 vocab / merges / special_tokens text : str use_cache : bool，False 时用一个"永远存不进东西"的假缓存来关闭缓存 repeats : int，重复轮数（>= 1），轮数越多计时越稳定 【返回】 float，字节/秒 """ nbytes = len(text.encode("utf-8")) total = 0.0 for _ in range(repeats): fresh = BPETokenizer(tok.vocab, tok.merges, tok.special_tokens) if not use_cache: fresh._cache = _NoCache() t0 = time.perf_counter() fresh.encode(text) total += time.perf_counter() - t0 return nbytes * repeats / total def demo_speed(tok): """演示 7：编码速度与缓存的作用。""" print_header("演示 7：编码吞吐量") text = CORPUS_PATH.read_text(encoding="utf-8") speed_cached = measure_throughput(tok, text, use_cache=True) speed_nocache = measure_throughput(tok, text, use_cache=False) print(f" 有缓存：{speed_cached / 1e6:.2f} MB/秒") print(f" 无缓存：{speed_nocache / 1e6:.2f} MB/秒 ← 缓存让速度提高了约 {speed_cached / speed_nocache:.1f} 倍") pile_bytes = 825e9 days = pile_bytes / speed_cached / 86400 print(f"\n 按有缓存的速度，单进程编码 The Pile（825 GB）约需 {days:.1f} 天。") print(" 而且这是乐观估计：我们的词表只有 500、语料用词高度重复（缓存命中率很高）。") print(" → 真实系统会用多进程并行，或者用 Rust 等编译型语言实现（tiktoken、HF tokenizers）。") class _NoCache(dict): """一个永远存不进东西的字典，用来关闭缓存做对比实验。""" def __setitem__(self, key, value): pass def main(): demo_pretokenize() demo_special_tokens() tok = demo_train() demo_encode_decode(tok) demo_encode_iterable(tok) demo_save_load(tok) demo_speed(tok) print("\n全部演示运行完毕。05_tokenizer_experiments.py 会用这个分词器做更多实验。") if __name__ == "__main__": main()