minbpe

用于大型语言模型分词中常用的字节对编码（BPE）算法的极简、整洁代码。BPE算法之所以称为“字节级”，是因为它直接作用于UTF-8编码的字符串。

这一算法由OpenAI的GPT-2论文及其配套的GPT-2代码发布推广至大型语言模型领域。Sennrich等人，2015年被引用为在自然语言处理应用中使用BPE的原始参考文献。如今，所有现代大型语言模型（例如GPT、Llama、Mistral）都采用该算法来训练其分词器。

本仓库包含两种分词器，它们都能实现分词器的三大基本功能：1）基于给定文本训练分词器词汇表和合并规则；2）将文本编码为标记；3）将标记解码回文本。仓库中的文件如下：

1. minbpe/base.py：实现了Tokenizer类，这是基础类。它包含了train、encode和decode的空实现，以及保存/加载功能，还有一些常用工具函数。这个类并不打算直接使用，而是作为其他类的基类来继承。
2. minbpe/basic.py：实现了BasicTokenizer，这是最简单的BPE算法实现，直接作用于文本。
3. minbpe/regex.py：实现了RegexTokenizer，它进一步通过正则表达式模式分割输入文本，这一步属于预处理阶段，先按类别（如字母、数字、标点符号）对输入文本进行拆分，然后再进行分词。这样可以确保不会跨类别边界发生合并。这一方法最初见于GPT-2论文，至今仍在GPT-4中沿用。该类还负责处理特殊标记（如有）。
4. minbpe/gpt4.py：实现了GPT4Tokenizer。这个类是对RegexTokenizer（上文第2项）的轻量封装，完全复现了tiktoken库中GPT-4的分词方式。封装处理了一些细节问题，比如如何精确恢复分词器中的合并规则，以及一些不幸的（可能是历史遗留的）单字节标记排列问题。

最后，脚本train.py会基于输入文本tests/taylorswift.txt（这是她的维基百科词条，也是个梗）训练这两种主要分词器，并将词汇表保存到磁盘以便可视化。在我的(M1)MacBook上，这个脚本大约运行25秒。

以上所有文件都非常简短且注释详尽，每份文件底部还附有使用示例。

快速入门

作为最简单的例子，我们可以重现维基百科关于BPE的文章，如下所示：

from minbpe import BasicTokenizer
tokenizer = BasicTokenizer()
text = "aaabdaaabac"
tokenizer.train(text, 256 + 3) # 256是字节标记，然后做3次合并
print(tokenizer.encode(text))
# [258, 100, 258, 97, 99]
print(tokenizer.decode([258, 100, 258, 97, 99]))
# aaabdaaabac
tokenizer.save("toy")
# 写入两个文件：toy.model（用于加载）和toy.vocab（用于查看）

根据维基百科，对输入字符串“aaabdaaabac”进行3次合并后，得到的结果是“XdXac”，其中X=ZY，Y=ab，Z=aa。需要注意的是，minbpe总是先分配256个单独的字节作为标记，然后根据需要再进行字节合并。因此，对于a=97，b=98，c=99，d=100（它们的ASCII值）。当(a,a)合并为Z时，Z会变成256。同样，Y会变成257，X会变成258。所以我们从256个字节开始，经过3次合并得到上述结果，预期输出为[258, 100, 258, 97, 99]。

推理：与GPT-4对比

我们可以验证RegexTokenizer与tiktoken中的GPT-4分词器功能一致，如下所示：

text = "hello123!!!? (안녕하세요!) 😉"

# tiktoken
import tiktoken
enc = tiktoken.get_encoding("cl100k_base")
print(enc.encode(text))
# [15339, 4513, 12340, 30, 320, 31495, 230, 75265, 243, 92245, 16715, 57037]

# 我们的
from minbpe import GPT4Tokenizer
tokenizer = GPT4Tokenizer()
print(tokenizer.encode(text))
# [15339, 4513, 12340, 30, 320, 31495, 230, 75265, 243, 92245, 16715, 57037]

（你需要先pip install tiktoken才能运行）。在底层，GPT4Tokenizer只是对RegexTokenizer的轻量封装，传入了GPT-4的合并规则和特殊标记。我们还可以确保特殊标记被正确处理：

text = "<|endoftext|>hello world"

# tiktoken
import tiktoken
enc = tiktoken.get_encoding("cl100k_base")
print(enc.encode(text, allowed_special="all"))
# [100257, 15339, 1917]

# 我们的
from minbpe import GPT4Tokenizer
tokenizer = GPT4Tokenizer()
print(tokenizer.encode(text, allowed_special="all"))
# [100257, 15339, 1917]

注意，和tiktoken一样，我们必须在调用encode时显式声明要使用和解析特殊标记。否则可能会成为重大隐患，无意中用特殊标记对攻击者控制的数据（例如用户提示）进行分词。allowed_special参数可以设置为“all”、“none”，或者允许使用的特殊标记列表。

训练

与tiktoken不同，这段代码允许你训练自己的分词器。原则上，据我所知，如果你用一个包含10万个词汇的大数据集来训练RegexTokenizer，就能重现GPT-4的分词器。

有两种路径可供选择。首先，如果你不想处理正则表达式模式对文本进行拆分和预处理的复杂性，也不在意特殊标记，那么可以选用BasicTokenizer。你可以训练它，然后像下面这样进行编码和解码：

from minbpe import BasicTokenizer
tokenizer = BasicTokenizer()
tokenizer.train(very_long_training_string, vocab_size=4096)
tokenizer.encode("hello world") # 字符串 -> tokens
tokenizer.decode([1000, 2000, 3000]) # tokens -> 字符串
tokenizer.save("mymodel") # 保存模型文件mymodel.model和词汇表mymodel.vocab
tokenizer.load("mymodel.model") # 从磁盘加载模型，词汇表仅用于可视化

如果你希望模仿OpenAI文本分词器的做法，采用正则表达式模式按类别拆分文本是个不错的选择。GPT-4的模式是RegexTokenizer的默认设置，因此你可以简单地这样做：

from minbpe import RegexTokenizer
tokenizer = RegexTokenizer()
tokenizer.train(very_long_training_string, vocab_size=32768)
tokenizer.encode("hello world") # 字符串 -> tokens
tokenizer.decode([1000, 2000, 3000]) # tokens -> 字符串
tokenizer.save("tok32k") # 保存模型文件tok32k.model和词汇表tok32k.vocab
tokenizer.load("tok32k.model") # 从磁盘加载模型

当然，你需要根据数据集的大小调整词汇表的大小。

特殊标记。最后，你可能希望在分词器中添加一些特殊标记。使用register_special_tokens函数注册这些标记。例如，如果你用vocab_size=32768训练，前256个标记是原始字节标记，接下来的32768-256个是合并标记，之后你可以添加特殊标记。最后一个“真正”的合并标记的ID是32767（vocab_size - 1），所以你的第一个特殊标记应该紧接其后，ID正好是32768。因此：

from minbpe import RegexTokenizer
tokenizer = RegexTokenizer()
tokenizer.train(very_long_training_string, vocab_size=32768)
tokenizer.register_special_tokens({"<|endoftext|>": 32768})
tokenizer.encode("<|endoftext|>hello world", allowed_special="all")

当然，你还可以根据需要继续添加更多标记。最后，我想强调一下，我尽力让代码本身保持干净、易读且可扩展。你不应该害怕阅读代码并理解它的运作原理。测试用例也是寻找更多使用示例的好地方。说到这里：

测试

我们使用pytest库进行测试。所有测试都位于tests/目录下。如果你还没有安装pytest，先执行pip install pytest，然后：

$ pytest -v .

运行测试。（-v表示详细输出，看起来更美观）

社区扩展

• gnp/minbpe-rs：minbpe的Rust实现，与Python版本几乎一一对应

练习

对于那些想学习BPE的人，这里提供了一个循序渐进的练习，教你如何一步步构建自己的minbpe。请参阅exercise.md。

讲座

我在本仓库中的代码是在这个YouTube视频中讲解的。你也可以在lecture.md中找到这篇讲座的文字版。

待办事项

• 编写一个更优化的Python版本，能够处理大文件和大规模词汇表
• 编写一个更优化的C或Rust版本（仔细思考）
• 将GPT4Tokenizer更名为GPTTokenizer，并支持GPT-2/GPT-3/GPT-3.5？
• 编写一个类似于GPT4Tokenizer的LlamaTokenizer（即尝试实现sentencepiece的等效功能）

许可证

MIT

minbpe 快速上手指南

环境准备

• Python 3.7+
• pip
• git（通常已预装）

安装步骤

pip install git+https://github.com/karpathy/minbpe -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

from minbpe import BasicTokenizer
tokenizer = BasicTokenizer()
text = "aaabdaaabac"
tokenizer.train(text, 256 + 3)  # 256字节令牌 + 3次合并
print(tokenizer.encode(text))  # [258, 100, 258, 97, 99]
print(tokenizer.decode([258, 100, 258, 97, 99]))  # aaabdaaabac
tokenizer.save("toy")  # 生成 toy.model 和 toy.vocab