# 用Python从零构建BPE分词器：深入算法核心与工程实践 如果你最近在玩大语言模型，大概率会频繁接触到一个词：**Tokenizer**，也就是分词器。无论是调用OpenAI的API，还是本地部署Llama、ChatGLM，输入的文本都需要先经过这个“翻译官”的转换，变成模型能理解的数字序列。而**BPE（Byte Pair Encoding）**，正是当今众多顶尖模型分词器背后的核心算法。理解BPE，不仅是理解现代NLP的基石，更是你从“调包侠”迈向“造轮子”工程师的关键一步。 网上关于BPE原理的科普文章很多，但大多停留在概念描述，或者直接甩给你一个Hugging Face的`AutoTokenizer`调用示例。这就像只告诉你汽车能跑，却不给你看发动机。今天，我们换个玩法：**不用任何现成的分词库，只用纯Python和标准库，从零开始，手搓一个功能完整的BPE分词器**。我们会深入每一个细节，从最基础的词频统计，到合并规则的贪婪选择，再到处理多语言文本的UTF-8编码技巧，最后实现与Hugging Face生态兼容的编码解码接口。这个过程不仅能让你彻底吃透BPE，更能锻炼你解决实际工程问题的能力。 ## 1. 理解BPE：为什么是它统治了LLM时代？ 在深入代码之前，我们有必要先厘清一个根本问题：在众多分词算法中，为什么BPE能脱颖而出，成为GPT系列、Llama等主流大模型的选择？ 传统的分词方法无非两种极端：按词（Word）分和按字符（Char）分。按词分，词表会爆炸（想想英文的时态变形、中文的海量词汇），遇到生僻词（OOV）就直接抓瞎。按字符分，词表是小了（英文26个字母，中文几千常用字），但序列长度剧增，且单个字符的语义信息太稀薄，模型学习效率低下。 **BPE的精妙之处，在于找到了一条“中庸之道”**。它从字符级别开始，通过不断合并最高频的相邻符号对，像搭积木一样，逐步构建出子词（Subword）词表。高频词（如“the”、“ing”）会保持完整，低频词或生僻词则被拆分成更常见的子词组合。例如，“unfortunately”可能被拆成“un”、“for”、“tun”、“ate”、“ly”。这完美平衡了词表大小与语义表达力。 但BPE也有其局限性。它本质上是一种**贪婪算法**，每次只合并当前频次最高的符号对，这个选择是局部最优的，但不一定是全局最优。这可能导致同一个词在不同上下文中被拆分成不同的子词序列，带来一定的歧义。不过，在巨大的训练数据量面前，这种歧义带来的噪声被证明是模型可以学习和容忍的。 另一个关键演进是**BBPE（Byte-level BPE）**，由GPT-2引入并广泛用于后续模型。它的核心突破是将合并的基本单元从“字符”变成了“字节”。初始词表就是256个可能的字节值。这样做的好处是**彻底解决了OOV问题**，因为任何文本都能用这256个字节表示。同时，它天然支持多语言混合文本，因为所有语言在底层都是字节流。代价是，对于像中文这样单个字符需要多个字节编码的语言，序列长度会比字符级BPE更长。 下面的表格对比了几种主流Subword算法的核心思想与典型应用： | 算法 | 核心思想 | 典型应用模型 | 主要特点 | | :--- | :--- | :--- | :--- | | **BPE** | 从字符开始，贪婪合并最高频相邻符号对 | GPT-1, RoBERTa | 简单有效，平衡词表与效率，但依赖预分词。 | | **BBPE** | 在字节级别进行BPE合并 | GPT-2, Llama, Bloom | 词表极小(256基础)，彻底无OOV，支持多语言，但序列可能更长。 | | **WordPiece** | 合并能最大提升语言模型概率的相邻子词 | BERT, DistilBERT | 类似BPE，但合并准则基于概率（互信息），对拼写错误敏感。 | | **Unigram** | 从一个超大词表开始，根据语言模型概率迭代丢弃子词 | SentencePiece (可选) | 能输出概率化分词结果，更灵活，但训练更复杂。 | | **SentencePiece** | 工具包，将输入视为Unicode序列，支持BPE/Unigram等算法 | XLNet, T5 | 无需预分词，直接处理原始文本，空白字符作为显式标记处理。 | > 注意：BPE和WordPiece通常需要一个预分词（Pre-tokenization）步骤，比如用空格初步分割英文单词。而SentencePiece和BBPE可以直接处理原始字节流，这是它们在工程上的一个重要区别。 理解了这些背景，我们就可以动手了。我们的目标是构建一个**字符级BPE分词器**，并在这个过程中，让你理解如何扩展到字节级（BBPE），以及如何与现有生态对接。 ## 2. 搭建基础骨架：词频统计与合并循环 让我们从最核心的算法流程开始。BPE训练的本质是一个迭代合并的过程，我们需要几个基础函数来支撑这个循环。 首先，我们需要从语料中构建初始词汇表。这里假设我们的输入已经是经过预分词（比如按空格分割单词）并添加了词尾标记的序列。 ```python import re from collections import defaultdict, Counter from typing import Dict, List, Tuple def get_stats(vocab: Dict[str, int]) -> Dict[Tuple[str, str], int]: """ 统计当前词汇表中所有相邻符号对的出现频率。 参数: vocab: 字典，键为单词（符号用空格分隔），值为该单词的频次。 例如：{'l o w </w>': 5, 'n e w e s t </w>': 6} 返回: 一个字典，键为符号对 (symbol_i, symbol_i+1)，值为该符号对的总频次。 """ pairs = defaultdict(int) for word, freq in vocab.items(): symbols = word.split() # 将单词分割成符号列表 for i in range(len(symbols) - 1): pair = (symbols[i], symbols[i+1]) pairs[pair] += freq return pairs def merge_vocab(pair: Tuple[str, str], vocab: Dict[str, int]) -> Dict[str, int]: """ 将词汇表中所有出现的指定符号对进行合并。 参数: pair: 需要合并的符号对，如 ('e', 's')。 vocab: 当前的词汇表。 返回: 合并指定符号对后的新词汇表。 """ first, second = pair # 创建正则表达式模式，匹配完整的符号对（前后有空格或边界） # 使用re.escape防止特殊字符被解释 pattern = re.compile(r'(?<!\S)' + re.escape(f'{first} {second}') + r'(?!\S)') new_vocab = {} for word in vocab: # 用合并后的新符号替换所有匹配的符号对 new_word = pattern.sub(f'{first}{second}', word) new_vocab[new_word] = vocab[word] return new_vocab ``` 有了这两个核心函数，BPE训练的主循环就非常直观了。我们设定一个目标词表大小 `vocab_size`，它等于基础字符数加上我们想要通过合并新增的子词数量。 ```python def train_bpe(raw_text: str, vocab_size: int = 50) -> Tuple[Dict[str, int], List[Tuple[str, str]]]: """ BPE训练函数。 参数: raw_text: 原始文本字符串。 vocab_size: 期望的最终词表大小（包含基础字符）。 返回: final_vocab: 训练结束后的词汇表（符号->频次）。 merges: 按顺序记录的所有合并操作对。 """ # 1. 预分词与添加词尾标记 words = raw_text.lower().split() # 简单按空格分割并转为小写 vocab = Counter([f"{' '.join(list(word))} </w>" for word in words]) # 2. 初始化基础词表（所有唯一字符） base_tokens = set() for word in vocab: base_tokens.update(word.split()) # 移除词尾标记</w>，它本身就是一个特殊符号 base_tokens.discard('</w>') # 我们最终想要的词表是基础字符 + 合并产生的新子词 # 计算需要合并的次数 num_merges = vocab_size - len(base_tokens) - 1 # -1 给 </w> merges = [] # 记录合并历史 current_vocab = dict(vocab) for i in range(num_merges): pairs = get_stats(current_vocab) if not pairs: break # 没有更多可以合并的对了 # 找到频率最高的符号对 best_pair = max(pairs, key=pairs.get) best_freq = pairs[best_pair] if best_freq < 2: # 可选：如果最高频次小于2，提前停止 break # 执行合并 current_vocab = merge_vocab(best_pair, current_vocab) merges.append(best_pair) # 可选：打印每次合并的信息 # print(f"Iteration {i+1}: Merge {best_pair} (freq: {best_freq})") # 3. 构建最终的符号到ID的映射 # 首先收集所有出现的符号 all_tokens = set() for word in current_vocab: all_tokens.update(word.split()) # 给符号排序，通常给高频词或特殊符号靠前的ID # 这里简单按符号字符串排序，实践中可能按频率排序 sorted_tokens = sorted(all_tokens) # 确保</w>在词表中 if '</w>' not in sorted_tokens: sorted_tokens.append('</w>') token_to_id = {token: idx for idx, token in enumerate(sorted_tokens)} # 我们还需要一个从合并对到排序ID的映射，用于编码 # 但更简单的方法是，我们最终使用一个基于`merges`顺序的编码函数 return current_vocab, merges ``` 让我们用一个经典的迷你语料库来测试一下这个训练过程： ```python if __name__ == "__main__": # 一个简单的测试语料 corpus = "low lower newest widest low low low low lower newest newest newest newest newest widest widest widest" vocab, merge_ops = train_bpe(corpus, vocab_size=30) print("最终词汇表（部分）:") for word, freq in list(vocab.items())[:10]: print(f" {word}: {freq}") print(f"\n合并操作记录 (共{len(merge_ops)}次):") for i, pair in enumerate(merge_ops): print(f" {i+1}: {pair[0]} + {pair[1]} -> {pair[0]+pair[1]}") ``` 运行这段代码，你会看到算法如何一步步地将“e s”合并成“es”，再将“es t”合并成“est”，最后甚至可能产生“est</w>”这样的子词。这个过程直观地展示了BPE如何从数据中自动学习常见的字符组合。 ## 3. 实现编码与解码：从文本到ID序列 训练得到`merges`列表后，我们如何对一个新句子进行分词（编码）？编码过程需要模拟训练时的合并操作，但方向相反：对于一个新单词，我们从其字符序列开始，然后按照`merges`记录的顺序，尝试应用每一个合并规则。 这里有一个关键点：合并必须是最长匹配优先，并且要遵循训练时的顺序。下面是一个直接但清晰的编码实现： ```python class SimpleBPETokenizer: def __init__(self, merges: List[Tuple[str, str]]): """ 初始化BPE分词器。 参数: merges: 训练得到的合并操作列表。 """ self.merges = merges # 根据merges推导出词表 self.vocab = self._build_vocab_from_merges() self.token_to_id = {token: i for i, token in enumerate(self.vocab)} self.id_to_token = {i: token for token, i in self.token_to_id.items()} def _build_vocab_from_merges(self) -> List[str]: """从合并历史中重建词表。""" # 基础词表是所有出现在merges中的字符 base_chars = set() for a, b in self.merges: base_chars.update(a) base_chars.update(b) # 加上词尾标记 vocab_set = set(base_chars) vocab_set.add('</w>') # 应用所有合并，生成新的子词 # 注意：这是一个简化版本，实际需要更严谨地模拟合并过程 # 这里我们直接假设合并产生的新符号就是 a+b for a, b in self.merges: new_token = a + b vocab_set.add(new_token) # 还需要考虑新token可能进一步参与合并，这里简化处理 # 一个完整的实现需要维护一个不断增长的符号集合 # 将集合转为列表并排序（排序策略会影响ID分配） return sorted(vocab_set) def encode_word(self, word: str) -> List[str]: """ 编码单个单词。 参数: word: 输入单词，如"lower"。 返回: 子词token列表，如['low', 'er</w>']。 """ # 初始化为字符列表，并添加词尾标记 tokens = list(word) + ['</w>'] # 遍历所有合并规则 for merge_first, merge_second in self.merges: new_token = merge_first + merge_second i = 0 while i < len(tokens) - 1: # 如果找到匹配的相邻对，则合并 if tokens[i] == merge_first and tokens[i+1] == merge_second: tokens[i : i+2] = [new_token] # 合并后不移动i，因为新token可能继续与后面的合并 else: i += 1 return tokens def encode(self, text: str) -> List[int]: """ 编码完整文本。 参数: text: 输入文本字符串。 返回: token ID列表。 """ # 简单按空格预分词 words = text.lower().split() token_ids = [] for word in words: subword_tokens = self.encode_word(word) for token in subword_tokens: # 如果token不在词表中，可以回退到字符级别或标记为UNK # 这里简化处理，假设所有token都在词表中 token_ids.append(self.token_to_id.get(token, self.token_to_id.get('<unk>', 0))) return token_ids def decode(self, token_ids: List[int]) -> str: """ 将token ID序列解码回文本。 参数: token_ids: token ID列表。 返回: 解码后的文本字符串。 """ tokens = [self.id_to_token.get(idx, '<unk>') for idx in token_ids] # 将token拼接起来 text = ''.join(tokens) # 处理词尾标记：将</w>替换为空格，并去除末尾可能多余的空格 text = text.replace('</w>', ' ') return text.strip() ``` 这个实现虽然直观，但在效率上存在问题。`encode_word`函数中的while循环在每次应用合并规则时都要遍历整个token列表，时间复杂度较高。在实际的工业级实现中（例如Hugging Face的tokenizers库），会采用更高效的算法，例如使用优先队列来管理所有可能的合并对。 让我们测试一下编码解码过程： ```python # 接续之前的训练代码 tokenizer = SimpleBPETokenizer(merge_ops) test_sentence = "the newest lower window" encoded_ids = tokenizer.encode(test_sentence) print(f"句子: '{test_sentence}'") print(f"编码ID: {encoded_ids}") print(f"对应Token: {[tokenizer.id_to_token[i] for i in encoded_ids]}") decoded_text = tokenizer.decode(encoded_ids) print(f"解码文本: '{decoded_text}'") ``` 你会看到像“newest”这样的词可能被拆分成“new”和“est</w>”，而“window”可能因为不在训练语料中，被拆分成更细的字符组合。这就是BPE处理OOV的能力。 ## 4. 处理多语言与UTF-8：迈向字节级BPE(BBPE) 我们的简单实现处理英文尚可，但面对中文、日文或混合文本时就会遇到问题。因为我们的基础单元是“字符”，而不同语言的字符集差异巨大。这就是**字节对编码（Byte-level BPE, BBPE）** 闪亮登场的时候。 BBPE的核心思想是将所有文本先编码为UTF-8字节序列，然后在**字节层面**进行BPE合并。初始词表就是0-255这256个字节值。这样做的好处是： 1. **词表极小且固定**：基础词表永远是256，与语言无关。 2. **绝对无OOV**：任何文本都能被表示为字节序列。 3. **多语言无缝支持**：统一了所有文字的表示空间。 让我们修改之前的代码，实现一个简化版的BBPE。关键变化在于预处理步骤： ```python def bytes_to_visible_string(byte_val: int) -> str: """将字节值转换为可显示的字符串表示，便于调试。""" if 32 <= byte_val <= 126: # 可打印ASCII范围 return chr(byte_val) else: return f"<0x{byte_val:02x}>" # 显示为十六进制，如 <0xe8> class ByteLevelBPETokenizer: def __init__(self, merges: List[Tuple[int, int]], vocab_size: int = 1000): """ 初始化字节级BPE分词器。 参数: merges: 合并操作列表，每个元素是一个字节值对 (int, int)。 vocab_size: 目标词表大小。 """ self.merges = merges # 例如 [(101, 115), (257, 100)]，其中257是合并后新token的虚拟ID self.vocab_size = vocab_size # 基础词表：256个字节 self.base_vocab = {i: bytes([i]) for i in range(256)} # 构建扩展词表 self.vocab = self._build_vocab() self.token_to_id = {token: idx for idx, token in enumerate(self.vocab)} # 我们需要一个从字节序列到token的映射，这里简化，实际编码需要模拟合并过程 # 更高效的做法是构建一个基于merges的前缀树（Trie） def _build_vocab(self) -> List[bytes]: """构建词表，包含基础字节和合并产生的新token。""" vocab_tokens = [bytes([i]) for i in range(256)] next_id = 256 # 模拟合并过程，生成新token # 注意：这是一个概念性实现。实际BBPE中，新token是字节序列的拼接 for first_byte, second_byte in self.merges: # 在实际BBPE中，我们需要找到first_byte和second_byte对应的token（可能是基础字节或多字节token） # 然后将其字节序列拼接起来，形成新token # 这里为了演示，我们假设first_byte和second_byte就是基础字节值 new_token = bytes([first_byte]) + bytes([second_byte]) vocab_tokens.append(new_token) next_id += 1 if len(vocab_tokens) >= self.vocab_size: break return vocab_tokens def encode_text(self, text: str) -> List[int]: """ 将文本编码为token ID序列。 简化版：先将文本转为UTF-8字节序列，然后尝试应用合并规则。 """ # 1. 转换为UTF-8字节 byte_sequence = list(text.encode('utf-8')) # 2. 模拟BPE合并过程（这里是非常简化的贪心匹配） # 实际实现需要更复杂的算法来处理重叠匹配和最长匹配 tokens = byte_sequence.copy() # 将merges按照长度（可能长度）或其他优先级排序 # 这里我们假设merges已经是按学习顺序排列的 for merge_first, merge_second in self.merges: i = 0 while i < len(tokens) - 1: # 注意：这里的tokens元素是整数（字节值） if tokens[i] == merge_first and tokens[i+1] == merge_second: # 合并：用一个新的ID（例如256+）替换这两个字节 # 在实际中，我们需要管理一个从字节序列到新ID的映射 # 这里简化，用占位符-1表示合并后的token tokens[i : i+2] = [-1] # 占位符 # 由于合并改变了序列长度和索引，这里逻辑需要更精细处理 # 为简化演示，我们跳出复杂循环 break else: i += 1 # 简化起见，只演示一次合并 break # 3. 将字节值映射到token ID（简化：直接映射到0-255） token_ids = [] for byte_val in byte_sequence[:10]: # 只取前10个字节演示 token_ids.append(byte_val) # 这里ID就是字节值本身 return token_ids def decode_ids(self, token_ids: List[int]) -> str: """将token ID解码回文本。""" byte_list = [] for tid in token_ids: if tid < 256: byte_list.append(tid) else: # 对于扩展词表中的token，需要查找其对应的字节序列 # 这里简化，假设扩展token也是直接可用的字节序列 # 在实际中，我们需要维护一个id到字节序列的映射 pass # 将字节列表转换为bytes对象，然后解码为字符串 byte_data = bytes(byte_list) try: return byte_data.decode('utf-8') except UnicodeDecodeError: # 如果字节序列不构成有效的UTF-8，可能因为只取了部分序列 return "[部分字节序列，无法完全解码]" ``` > 提示：上述BBPE实现是高度简化的概念演示。真实的BBPE实现（如Hugging Face的`GPT2Tokenizer`）要复杂得多，它需要高效地管理一个从字节序列到token ID的映射，并在编码时执行最长匹配。这通常通过构建一个**前缀树（Trie）** 来实现，树的边是字节，节点对应token ID。 让我们看一个BBPE如何处理中文的例子： ```python # 概念性演示 text_zh = "自然语言处理" byte_repr = list(text_zh.encode('utf-8')) print(f"文本: {text_zh}") print(f"UTF-8字节序列 (十六进制): {[hex(b) for b in byte_repr]}") print(f"字节序列长度: {len(byte_repr)}") print("说明：在BBPE中，这", len(byte_repr), "个字节将作为初始输入，进行BPE合并。") print("对于常见的中文字符组合，BBPE可能会学习到对应的多字节token，从而减少token数量。") ``` 输出会显示，一个简单的6汉字中文句子，在UTF-8编码下可能变成18个甚至更多的字节。BBPE的训练目标就是从这些字节序列中，找出跨语言的、高频的字节组合，形成共享的子词单元。 ## 5. 工程化与Hugging Face生态兼容 自己实现的玩具分词器用于学习很棒，但要投入实际使用，尤其是想和现有的Transformer模型（如加载Hugging Face模型）一起工作，我们需要让分词器符合一定的接口规范。Hugging Face的`transformers`库定义了一个通用的`PreTrainedTokenizer`基类。 我们的目标是创建一个自定义分词器，它可以被保存和加载，并且拥有`__call__`、`encode`、`decode`、`convert_tokens_to_ids`等标准方法。这里我们实现一个精简的、与我们的BPE逻辑兼容的类。 首先，我们需要一种持久化保存分词器状态（词表、合并规则、特殊标记等）和从磁盘加载的方式。通常，Hugging Face分词器使用一个`tokenizer.json`文件。 ```python import json import os class CustomBPETokenizer: def __init__(self, merges: List[Tuple[str, str]], vocab: Dict[str, int] = None, unk_token: str = "<unk>", pad_token: str = "<pad>", bos_token: str = "<s>", eos_token: str = "</s>"): """ 一个更工程化的BPE分词器。 参数: merges: 合并操作列表。 vocab: 可选的词汇表（token->频率）。如果为None，将从merges推导。 unk_token: 未知token。 pad_token: 填充token。 bos_token: 句子开始token。 eos_token: 句子结束token。 """ self.merges = merges self.unk_token = unk_token self.pad_token = pad_token self.bos_token = bos_token self.eos_token = eos_token # 构建词表 if vocab is not None: self.vocab = vocab else: self.vocab = self._build_vocab_from_merges() # 添加特殊token special_tokens = [unk_token, pad_token, bos_token, eos_token] for token in special_tokens: if token not in self.vocab: self.vocab[token] = 0 # 频率为0或一个固定值 # 创建双向映射 self._create_mappings() # 构建用于快速编码的前缀树（Trie）或合并规则映射（这里简化） self._build_encoder_cache() def _build_vocab_from_merges(self) -> Dict[str, int]: """从合并历史重建词汇表（带频率估算，这里简化设为1）。""" vocab_set = set() # 收集所有基础字符 for a, b in self.merges: vocab_set.update(a) vocab_set.update(b) vocab_set.add('</w>') # 应用合并规则生成新token # 这是一个简化模拟，实际频率需要从训练语料统计 for a, b in self.merges: vocab_set.add(a + b) # 为每个token分配一个虚拟频率（实际应从训练数据统计） return {token: 1 for token in vocab_set} def _create_mappings(self): """创建token到id和id到token的映射。""" # 按一定顺序排序token，例如先特殊token，再按频率或字母顺序 all_tokens = list(self.vocab.keys()) # 将特殊token放在前面 specials = [self.unk_token, self.pad_token, self.bos_token, self.eos_token, '</w>'] # 移除已存在的特殊token，然后按原顺序添加回来 for token in specials: if token in all_tokens: all_tokens.remove(token) sorted_tokens = specials + sorted([t for t in all_tokens if t not in specials]) self.token_to_id = {token: idx for idx, token in enumerate(sorted_tokens)} self.id_to_token = {idx: token for token, idx in self.token_to_id.items()} self.vocab_size = len(sorted_tokens) def _build_encoder_cache(self): """构建编码缓存，加速单词到token的转换。""" # 在实际实现中，这里会构建一个Trie树或合并规则的有序列表 # 用于在编码时快速找到最长匹配的子词 # 为简化，我们只存储merges，编码时使用稍慢的循环方法 self._merges_by_length = sorted(self.merges, key=lambda x: len(x[0]+x[1]), reverse=True) def tokenize(self, text: str) -> List[str]: """将文本分割成token列表（字符串形式）。""" words = text.lower().split() # 预分词，实际可能更复杂 tokens = [] for word in words: word_tokens = self._tokenize_word(word) tokens.extend(word_tokens) return tokens def _tokenize_word(self, word: str) -> List[str]: """分词单个单词。使用基于merges的贪婪最长匹配。""" # 添加词尾标记 chars = list(word) + ['</w>'] # 循环应用所有合并规则，直到无法合并 changed = True while changed and len(chars) > 1: changed = False # 尝试所有可能的合并对（按学习顺序或长度优先） for first, second in self.merges: i = 0 while i < len(chars) - 1: if chars[i] == first and chars[i+1] == second: # 合并 chars[i : i+2] = [first + second] changed = True # 合并后继续检查当前位置（因为新token可能与后面的字符再次合并） else: i += 1 if changed: # 本轮有合并，重新开始扫描，确保所有规则被重新应用 break return chars def encode(self, text: str, add_special_tokens: bool = True) -> List[int]: """将文本编码为ID列表。""" tokens = self.tokenize(text) if add_special_tokens: tokens = [self.bos_token] + tokens + [self.eos_token] token_ids = [] for token in tokens: token_ids.append(self.token_to_id.get(token, self.token_to_id[self.unk_token])) return token_ids def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str: """将ID列表解码回文本。""" tokens = [] for tid in token_ids: if tid < len(self.id_to_token): token = self.id_to_token[tid] if skip_special_tokens and token in [self.bos_token, self.eos_token, self.pad_token, self.unk_token]: continue tokens.append(token) else: if not skip_special_tokens: tokens.append(self.unk_token) # 拼接tokens text = ''.join(tokens) # 处理词尾标记：将</w>替换为空格 text = text.replace('</w>', ' ') # 合并多余空格 text = ' '.join(text.split()) return text def save_pretrained(self, save_directory: str): """将分词器保存到目录。""" os.makedirs(save_directory, exist_ok=True) # 保存配置 config = { "merges": self.merges, "unk_token": self.unk_token, "pad_token": self.pad_token, "bos_token": self.bos_token, "eos_token": self.eos_token, "vocab": self.vocab, # 保存完整的词汇频率信息 } with open(os.path.join(save_directory, "tokenizer_config.json"), 'w', encoding='utf-8') as f: json.dump(config, f, ensure_ascii=False, indent=2) # 保存词表映射（Hugging Face格式通常是一个txt文件，每行一个token） with open(os.path.join(save_directory, "vocab.txt"), 'w', encoding='utf-8') as f: for token, _ in sorted(self.token_to_id.items(), key=lambda x: x[1]): f.write(token + '\n') # 保存合并规则（BPE格式，每行用空格分隔两个符号） with open(os.path.join(save_directory, "merges.txt"), 'w', encoding='utf-8') as f: for first, second in self.merges: f.write(f"{first} {second}\n") @classmethod def from_pretrained(cls, save_directory: str): """从目录加载分词器。""" config_path = os.path.join(save_directory, "tokenizer_config.json") with open(config_path, 'r', encoding='utf-8') as f: config = json.load(f) # 注意：从json加载的merges是列表的列表，需要转换为元组 merges = [tuple(pair) for pair in config["merges"]] return cls( merges=merges, vocab=config.get("vocab"), unk_token=config["unk_token"], pad_token=config["pad_token"], bos_token=config["bos_token"], eos_token=config["eos_token"], ) ``` 现在，我们可以像使用Hugging Face分词器一样使用它： ```python # 假设我们已经有了训练好的merges # 这里用一个小例子演示 example_merges = [('e', 's'), ('es', 't'), ('l', 'o'), ('lo', 'w')] tokenizer = CustomBPETokenizer(merges=example_merges, vocab_size=50) text = "The newest lower window is wide." encoded_ids = tokenizer.encode(text, add_special_tokens=True) print(f"编码结果: {encoded_ids}") print(f"Token数量: {len(encoded_ids)}") decoded_text = tokenizer.decode(encoded_ids, skip_special_tokens=True) print(f"解码结果: '{decoded_text}'") # 保存和加载 tokenizer.save_pretrained("./my_bpe_tokenizer") loaded_tokenizer = CustomBPETokenizer.from_pretrained("./my_bpe_tokenizer") print(f"加载后的分词器词表大小: {loaded_tokenizer.vocab_size}") ``` 这个自定义分词器已经具备了基本的功能。要将其完全集成到Hugging Face生态中，还需要继承`PreTrainedTokenizer`并实现更多抽象方法，但上面的框架已经揭示了核心逻辑。 ## 6. 性能优化与实战建议 自己实现BPE用于学习无妨，但在生产环境中，效率至关重要。以下是一些关键的优化方向和实战建议： **1. 编码算法优化** 我们之前实现的`_tokenize_word`函数效率是O(n*m)（n为单词长度，m为合并规则数）。对于包含数万甚至数十万合并规则的大词表，这是不可接受的。工业级实现采用以下策略： - **前缀树（Trie）**：将所有词汇表中的token（包括子词）构建成一棵前缀树。编码时，对输入单词进行最长匹配查找，时间复杂度接近O(n)。 - **缓存**：对常见的单词或子词的编码结果进行缓存，避免重复计算。 - **并行化**：对大批量文本进行编码时，可以利用多线程或多进程。 **2. 处理超大语料** 训练BPE需要遍历整个语料库多次统计频次。对于GB甚至TB级别的文本： - **流式处理**：不要一次性将全部语料加载到内存。可以分块读取，逐步更新符号对的频次统计。 - **近似计数**：使用布隆过滤器或Count-Min Sketch等概率数据结构来估计高频词对，减少内存占用。 - **分布式计算**：如果语料极大，可以考虑使用Spark或Dask进行分布式频次统计。 **3. 处理生僻字与未知语言** 即使采用BBPE，对于训练数据中极少出现的字符组合，编码效率依然很低（会被拆成很多单字节）。实践中可以： - **引入一个最低频次阈值**：在训练时，如果某个字节对的出现次数低于阈值，则停止合并，即使还没达到目标词表大小。 - **混合策略**：对于已知的高资源语言（如中英文），可以单独训练或使用现有的高质量词表；对于低资源语言，回退到字节级编码。 **4. 与现有模型兼容的注意事项** 如果你想用自己的分词器替换现有模型（如Llama）的分词器，需要注意： - **词表大小必须匹配**：模型的嵌入层（Embedding）大小等于词表大小。你不能随意改变词表大小，除非你也重新训练或调整模型嵌入层。 - **特殊token的ID**：`[PAD]`、`[UNK]`、`[BOS]`、`[EOS]`等特殊token的ID必须与模型预训练时一致，否则会导致行为异常。 - **分词粒度的影响**：改变分词粒度会影响模型看到的序列长度和语义单元。如果新分词器将同一个词拆成更多或更少的token，模型的性能可能会下降，因为其参数是基于原分词方式训练的。 一个实用的建议是，**优先使用模型原配的分词器**。只有在特定领域（如医学、法律文本）有大量未登录词，且经过实验证明原分词器严重影响性能时，才考虑训练领域自适应的BPE分词器，并通过持续预训练（Continual Pretraining）来让模型适应新的分词方式。 从头实现BPE的过程，就像亲手拆解一台精密的钟表。你看到了每个齿轮（函数）如何咬合，理解了为什么需要润滑油（优化技巧），也明白了它为何能在NLP的舞台上精准报时。这份理解，远比单纯调用`tokenizer.encode()`来得深刻。当你下次再遇到分词相关的诡异bug或性能瓶颈时，这份从零构建的经验，将成为你手中最可靠的调试指南和优化蓝图。