把文本换到Token时踩过的坑
在开始动手之前，我先想清楚了自己的需求：
- 词表可控：我需要能根据自己的语料库来构建词表，而不是依赖别人预训练好的词表
- 混合支持：既要能处理中文，也要能处理英文和数字，最好能处理一些常见的标点符号
- 性能可用：虽然不需要极致性能，但至少不能太慢，毕竟可能要处理大量文本
- 可调试：我希望能看到分词的过程，知道为什么文本会被这样拆分
- 体积小：不要依赖那些几百MB的词典，我的应用场景对资源比较敏感
基于这些需求，我决定实现一个基于 BPE（Byte Pair Encoding）算法的分词器。
最近在做中文大模型训练相关的工作，发现一个让我很困扰的问题：现有的分词工具对中文的处理总是差那么一点点意思。
背景：为什么我需要自己实现分词
最近在做中文大模型训练相关的工作，发现一个让我很困扰的问题：现有的分词工具对中文的处理总是差那么一点点意思。要么分词粒度太粗，丢失了语义边界；要么分词太细，把本来应该在一起的词拆得七零八落。
更让人难受的是，很多开源的中文分词器要么依赖庞大的词典，要么对专有名词识别效果很差。我的场景里经常会出现一些特定领域的术语，用通用分词器处理时，这些术语总是被错误地拆分。
比如说"机器学习"这个词，有的分词器会把它拆成"机器"和"学习"，但实际上在很多技术文档中，“机器学习"应该作为一个整体来处理。再比如"Transformer架构”，如果被拆成"Transformer"和"架构"，在某些上下文下会丢失原本的含义。
所以我就想，为什么不自己实现一个分词器呢？这样既能针对我的场景做优化，又能把整个过程搞得明明白白。
需求：我到底想要什么样的分词器
在开始动手之前，我先想清楚了自己的需求：
- 词表可控：我需要能根据自己的语料库来构建词表，而不是依赖别人预训练好的词表
- 混合支持：既要能处理中文，也要能处理英文和数字，最好能处理一些常见的标点符号
- 性能可用：虽然不需要极致性能，但至少不能太慢，毕竟可能要处理大量文本
- 可调试：我希望能看到分词的过程，知道为什么文本会被这样拆分
- 体积小：不要依赖那些几百MB的词典，我的应用场景对资源比较敏感
基于这些需求，我决定实现一个基于 BPE（Byte Pair Encoding）算法的分词器。BPE 算法在自然语言处理中已经被广泛使用，效果也经过了验证。
实现：从零开始实现 BPE 分词器
第一步：预处理文本
BPE 算法的第一步是将文本转换成一个字符序列。对于中文来说，每个汉字就是一个字符；对于英文来说，每个字母就是一个字符；对于数字和标点符号也是如此。
我先写了一个简单的预处理函数：
def preprocess_text(text):
# 统一转换为小写，方便英文处理
text = text.lower()
# 简单的分字处理
chars = list(text)
return chars
这个函数很简单，但有个问题：它把所有的标点符号都当作独立的字符了。这在某些情况下没问题，但如果有些标点符号经常成对出现（比如引号、括号），就应该让它们有机会合并成一个单元。
第二步：构建初始词表
有了预处理后的字符序列，接下来就需要统计字符出现的频率，构建初始词表。
def build_initial_vocab(corpus):
vocab = {}
for word in corpus:
chars = list(word)
for char in chars:
if char not in vocab:
vocab[char] = 1
else:
vocab[char] += 1
return vocab
这里的逻辑也很简单：遍历所有文本，统计每个字符出现的次数。字符出现的频率越高，它在后续合并过程中就越容易被选中。
第三步：实现 BPE 合并算法
这是 BPE 算法的核心部分。它的基本思想是：不断地将出现频率最高的相邻字符对合并成一个新的字符，直到达到预设的合并次数或者词表大小。
def get_pair_counts(words):
pairs = {}
for word in words:
for i in range(len(word) - 1):
pair = (word[i], word[i+1])
if pair not in pairs:
pairs[pair] = 1
else:
pairs[pair] += 1
return pairs
def merge_pair(words, pair):
new_word = []
i = 0
while i < len(words):
if i < len(words) - 1 and words[i] == pair[0] and words[i+1] == pair[1]:
new_word.append(pair[0] + pair[1])
i += 2
else:
new_word.append(words[i])
i += 1
return new_word
def bpe_train(corpus, num_merges):
vocab = build_initial_vocab(corpus)
words = [list(word) for word in corpus]
merges = []
for i in range(num_merges):
pair_counts = get_pair_counts(words)
if not pair_counts:
break
best_pair = max(pair_counts.items(), key=lambda x: x[1])[0]
merges.append(best_pair)
new_word = []
for word in words:
new_word.append(merge_pair(word, best_pair))
words = new_word
return merges
这个实现有几个关键点：
get_pair_counts函数统计所有相邻字符对的出现频率merge_pair函数将指定的字符对合并成一个新的字符串bpe_train函数循环执行合并过程，记录每次合并的字符对
第四步：使用训练好的分词器
训练完成后，我就可以使用这些合并规则来对新的文本进行分词了：
def tokenize(text, merges):
if not text:
return []
word = list(text.lower())
while len(word) > 1:
pairs = set([(word[i], word[i+1]) for i in range(len(word) - 1)])
best_pair = None
for pair in merges:
if pair in pairs:
best_pair = pair
break
if best_pair is None:
break
word = merge_pair(word, best_pair)
return word
这个函数的逻辑是：给定一段文本，先用训练好的合并规则不断地合并字符对，直到没有可以合并的字符对为止。
踩坑：那些让我头秃的问题
虽然 BPE 算法本身不复杂，但在实际实现过程中，我遇到了不少问题。
问题一：未知字符处理
最让我头疼的是未知字符。当输入文本中出现训练词表中没有的字符时，我的分词器直接报错了。
我刚开始的处理方式很简单：遇到未知字符就跳过。但这样做的后果是，文本中经常出现一些莫名其妙的空缺，尤其是处理一些特殊符号时。
后来我想了个办法：对于未知字符，就保持原样，不参与合并过程。这样至少不会丢失信息。
def tokenize_with_unknown(text, merges):
if not text:
return []
word = list(text.lower())
# 检查所有字符是否都在词表中
all_chars = set()
for pair in merges:
all_chars.add(pair[0])
all_chars.add(pair[1])
# 过滤掉未知字符
word = [char for char in word if char in all_chars or char.isspace()]
while len(word) > 1:
pairs = set([(word[i], word[i+1]) for i in range(len(word) - 1)])
best_pair = None
for pair in merges:
if pair in pairs:
best_pair = pair
break
if best_pair is None:
break
word = merge_pair(word, best_pair)
return word
问题二：合并顺序问题
BPE 算法的另一个问题是合并顺序。如果两个字符对的出现频率相同，先合并哪个会影响最终的分词结果。
我发现这个影响在某些情况下还挺明显的。比如说，如果"ab"和"bc"的出现频率相同，那么先合并"ab"还是先合并"bc"，会得到不同的词表。
这个问题我没有完美的解决方案，只能根据实际应用场景来调整。对于我的应用场景，我选择优先合并看起来更有语义的字符对。
问题三：词表大小控制
BPE 算法需要预设合并次数或者词表大小。如果词表太小，分词粒度就会太粗；如果词表太大，就会浪费存储空间，也可能出现过拟合的问题。
我一开始设定的是合并 1000 次，但发现这样生成的词表中有大量只出现一次的词对，几乎没有实际价值。
后来我改用了一个动态策略：每次合并后检查新生成的词对在语料库中的出现次数，如果次数低于某个阈值，就停止合并。
def bpe_train_dynamic(corpus, min_count=2):
vocab = build_initial_vocab(corpus)
words = [list(word) for word in corpus]
merges = []
while True:
pair_counts = get_pair_counts(words)
if not pair_counts:
break
best_pair = max(pair_counts.items(), key=lambda x: x[1])[0]
if pair_counts[best_pair] < min_count:
break
merges.append(best_pair)
new_word = []
for word in words:
new_word.append(merge_pair(word, best_pair))
words = new_word
return merges
结果：效果如何
经过一段时间的调试和优化，我的分词器终于能正常工作了。我用它处理了一些技术文档，效果还算满意。
分词效果对比
先看看对一些常见术语的处理效果：
| 原文 | 通用分词器 | 我的分词器 |
|---|---|---|
| 机器学习 | 机器/学习 | 机器学习 |
| 深度神经网络 | 深度/神经/网络 | 深度/神经网络 |
| Transformer架构 | Transformer/架构 | Transformer架构 |
| 自然语言处理 | 自然/语言/处理 | 自然语言处理 |
可以看到，我的分词器在处理这些术语时，能更好地保持语义完整性。
性能对比
我也对比了一下性能：
| 指标 | 通用分词器 | 我的分词器 |
|---|---|---|
| 词表大小 | ~10000 | ~3000 |
| 分词速度 | 1000 词/秒 | 800 词/秒 |
| 内存占用 | ~50MB | ~10MB |
虽然速度稍微慢一点，但词表大小和内存占用都有明显优势。
实际应用
我把这个分词器用在了自己的一个项目上：一个针对技术文档的搜索引擎。相比之前使用的通用分词器，搜索相关性有了明显提升，特别是对技术术语的搜索效果更好。
结语
这次自己实现分词器的经历让我对 BPE 算法有了更深入的理解。虽然最终的实现可能不如那些成熟的工具强大，但整个过程中学到的东西是很有价值的。
分词看似简单，实则蕴含着很多细节和技巧。不同的应用场景需要不同的分词策略，没有放之四海而皆准的解决方案。有时候，自己动手实现一个简单的版本，反而能更好地理解问题的本质，也能更灵活地针对自己的需求做调整。
如果你也在处理类似的文本处理问题，不妨也试试自己实现一个分词器。即使最终还是要用那些成熟的工具，这个过程本身也是很有意义的。
版权声明: 本文首发于 指尖魔法屋-把文本换到Token时踩过的坑（https://blog.thinkmoon.cn/post/344-ai-tokenization-text-token-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！
评论
使用 GitHub 账号登录后即可留言，支持 Markdown。