从零实现BPE分词器：为小型GPT上下文窗口生成更短的序列

字符级分词器是让小型 GPT 运行的最简单方法：其词汇量极小，不存在未知内容，且每一步都清晰透明。但代价是序列长度，因为每个字符都会占用上下文窗口中的一个位置。本指南用纯 Python 从头实现的字节对编码（BPE）来替代它。你将使用 WikiText-2 对其进行训练，以 JSON 格式保存学到的词汇表和合并规则，为模型生成指纹，重新编码数据集，从而清楚地了解使用子词标记重新训练模型时会发生哪些变化。

该代码假定已有小型 PyTorch GPT 模型以及处理过的 WikiText-2 文件。如果训练仍出现问题，需先稳定模型；可参考针对小型 PyTorch GPT 中各类故障的隔离测试。

为何字符标记已不够用

字符分词器会为包括空格在内的每个字符分配一个独立的标记。以这个简短短语为例：

artificial intelligence

它会将每个字符转换为一个标记：

['a', 'r', 't', 'i', 'f', 'i', 'c', 'i', 'a', 'l',
 ' ', 'i', 'n', 't', 'e', 'l', 'l', 'i', 'g', 'e',
 'n', 'c', 'e']

两个词对应23个标记。而子词分词器可以用更少的单元来表示相同的短语，例如：

['art', 'ificial', ' intelligence']

从23个标记变为3个标记。更短的序列意味着每个上下文窗口能容纳更多文本，训练所需的数据量更少，生成时的解码步骤也更少。该模型接收的是有意义的词块而非孤立的字母，常见的片段可被用作可重复使用的单元，而稀有词汇仍可通过更小的部分拼写出来。难点在于自动发现优质的子词单元，而这正是BPE的任务。

字符、单词与子词的对比

文本分割有三种主要方式，每种方式都会带来不同的权衡。

字符级分词

每个字符都是一个标记：

"learning"
→ ["l", "e", "a", "r", "n", "i", "n", "g"]

它的词汇量很小，没有未知标记，实现也极为简单。但作为交换，你得到的却是很长的序列、缓慢的生成速度，以及一个必须从单个字母组合单词的模型。

单词分词

每个单词都视为一个标记：

"machine learning works"
→ ["machine", "learning", "works"]

序列较短且易于阅读，但词汇量极为庞大，未见过的单词需要用特殊的标记表示，而诸如 "learn"、"learning" 和 "learned" 这类相关形式最终可能会变成三个彼此毫无关联的标识符。对于名称、拼写错误以及罕见单词的处理也很糟糕。

子词分词法

单词会被拆分成常见的片段：

"unbelievable"
→ ["un", "believ", "able"]

这种方式产生的序列长度短于字符，词汇量也小于单词。高频词仍可视为单个标记，低频词则会被分解为已知的组成部分，而前缀和后缀之类的片段可在多个单词中重复使用。唯一的实际代价是训练和实现该分词器更为复杂。

你可以将这三种方法视为一个谱系上的不同点：

Characters
← flexibility

Subwords
← balance

Words
← compactness

BPE如何学习其合并规则

BPE从最小的单元开始，即单个字符，然后反复将最常见的相邻字符对合并为一个新标记。以包含三个单词的极小语料库为例：

low
lower
lowest

最初，每个单词仅由其字符组成：

l o w
l o w e r
l o w e s t

接下来，列出出现的相邻符号对：

(l, o)
(o, w)
(w, e)
(e, r)
(e, s)
(s, t)

有一对格外突出：

(l, o)

它出现在这三个单词中，因此被合并为一个符号：

(l, o) → lo

在所有地方都替换之后，语料库看起来如下：

lo w
lo w e r
lo w e s t

再次统计后，出现频率最高的组合现在是：

(lo, w)

它也出现了三次，因此也会被合并：

(lo, w) → low

于是语料库变为：

low
low e r
low e s t

经过足够多的轮次后，出现频率最高的序列会变成独立的标记。

训练循环的输出结果

你即将实现的流程遵循以下步骤：

Training text
→ initialize character tokens
→ count adjacent token pairs
→ select most frequent pair
→ merge the selected pair
→ add merged token to vocabulary
→ repeat until target vocabulary size

它会生成两个输出结果，之后若要使用分词器就必须同时使用这两个结果。

词汇表

词汇表会将每个词元字符串映射为对应的整数ID，例如：

"a"       → 0
"b"       → 1
"th"      → 98
"ing"     → 127
" the"    → 204

合并规则

合并规则是在训练过程中被合并的成对数据按顺序列出的列表：

("t", "h") → "th"
("i", "n") → "in"
("in", "g") → "ing"
(" ", "the") → " the"

顺序至关重要。只有当创建 "in" 的前置规则已被应用后，后续的规则如 ("in", "g") 才能被触发。以不同顺序应用相同的规则会导致不同的分割结果。

在示例语料库上的首次实现

在开始处理 WikiText-2 之前，先在可以直观查看的字符串上练习该算法：

text = "low lower lowest"

将其拆分为字符后，得到的起始标记列表为：

tokens = list(text)
print(tokens)

打印后会显示所有字符，包括空格：

['l', 'o', 'w', ' ', 'l', 'o', 'w', 'e', 'r',
 ' ', 'l', 'o', 'w', 'e', 's', 't']

要找出可合并的候选项，需统计每一对相邻元素。将列表与自身偏移一个位置后的版本进行拼接即可得到所有相邻对，然后使用Counter来统计它们的数量：

from collections import Counter


def count_pairs(tokens):
    return Counter(
        zip(tokens[:-1], tokens[1:])
    )

检查最常见的配对：

pair_counts = count_pairs(tokens)
print(pair_counts.most_common(10))

输出结果大致如下：

[(('l', 'o'), 3),
 (('o', 'w'), 3),
 (('w', 'e'), 2),
 ...]

请注意，计数中也会出现 (' ', 'l')。由于空格是该方案中的普通字符，BPE 可以学习包含前导空格的标记，这就是后来出现诸如 " the" 这类标记的原因。

合并单个对

核心操作是从左到右遍历标记列表。每当当前标记与下一个标记匹配所选的对时，就会输出合并后的标记并跳过两个位置；否则则复制当前标记并向前移动一位。重叠部分会从左到右解决：在 a a a 以及 (a, a) 这对标记的情况下，只有前两个标记会合并。

def merge_pair(
    tokens,
    pair,
    merged_token,
):
    output = []
    index = 0

    while index < len(tokens):
        if (
            index < len(tokens) - 1
            and tokens[index] == pair[0]
            and tokens[index + 1] == pair[1]
        ):
            output.append(merged_token)
            index += 2
        else:
            output.append(tokens[index])
            index += 1

    return output

将其应用于合并 l 和 o：

tokens = merge_pair(
    tokens=tokens,
    pair=("l", "o"),
    merged_token="lo",
)

print(tokens)

结果比输入内容更短：

['lo', 'w', ' ', 'lo', 'w', 'e', 'r',
 ' ', 'lo', 'w', 'e', 's', 't']

每次合并都会删除该词对每出现一次对应的一个标记，这就是压缩效果的来源。

重复操作直至词汇表满

最简单的训练器只是循环执行以下操作：统计词对数量，选出出现频率最高的词对，在所有位置进行合并，记录新的标记及规则，一旦词汇表大小达到目标值则停止。如果不再存在词对，或者出现频率最高的词对出现次数少于两次，它也会提前停止，因为此时合并已无法产生任何可重复使用的内容。

def train_bpe_on_tokens(
    tokens,
    target_vocab_size,
):
    vocabulary = set(tokens)
    merges = []

    while len(vocabulary) < target_vocab_size:
        pair_counts = count_pairs(tokens)

        if not pair_counts:
            break

        best_pair, frequency = (
            pair_counts.most_common(1)[0]
        )

        if frequency < 2:
            break

        merged_token = "".join(best_pair)

        tokens = merge_pair(
            tokens=tokens,
            pair=best_pair,
            merged_token=merged_token,
        )

        vocabulary.add(merged_token)
        merges.append(best_pair)

        print(
            f"Merged {best_pair} "
            f"into {merged_token!r}, "
            f"frequency={frequency}"
        )

    return tokens, vocabulary, merges

在带有小型目标的简易字符串上运行它，便足以逐个观察到合并过程的发生：

final_tokens, vocabulary, merges = (
    train_bpe_on_tokens(
        tokens=list(text),
        target_vocab_size=20,
    )
)

这个版本有助于建立直观理解，但每次迭代都会重新扫描整个Python列表。在完整的WikiText语料库上，这种方式会变得非常慢，因此下面的实用版本改为在可配置的WikiText-2子集上进行训练。

设计可重用的分词器

分词器类应能够从原始文本的字符开始进行训练，实现编码与解码功能，过滤掉词汇表之外的字符，能够将数据保存到JSON文件中并从中加载，同时还能计算指纹值。围绕这个类，还有一个脚本会对每个WikiText-2分割后的文本重新编码，并显示BPE序列由此缩短了多少。

有一个设计决策需要明确说明。该分词器处理的是 Unicode 字符，而非原始字节。大型模型所使用的生产级分词器通常在字节层面工作，这意味着它们的基础词汇表由 256 种可能的字节值构成，任何输入都必然可以被表示。此处的基于字符的版本遵循相同的算法，且更易于理解，但它并不能直接替代那些分词器。

类结构

构造函数接收基础字符和有序合并列表。它从基础字符开始，依次添加每次合并的结果并剔除重复项，从而构建完整的词元列表。基于该列表，它会实现 stoi（字符串转 ID）和 itos（ID 转字符串）的查找功能，并记录词汇表的大小。由于 ID 来自列表中的位置，相同的基词元和合并操作始终会生成相同的 ID。

class BPETokenizer:
    def __init__(
        self,
        base_tokens,
        merges,
    ):
        self.base_tokens = list(
            base_tokens
        )

        self.merges = [
            tuple(pair)
            for pair in merges
        ]

        self.tokens = list(
            self.base_tokens
        )

        for left, right in self.merges:
            merged = left + right

            if merged not in self.tokens:
                self.tokens.append(merged)

        self.stoi = {
            token: index
            for index, token in enumerate(
                self.tokens
            )
        }

        self.itos = {
            index: token
            for token, index
            in self.stoi.items()
        }

        self.vocab_size = len(
            self.tokens
        )

因此，生成的词汇表由两部分组成：

Base character tokens
+
learned merged tokens

基于真实文本的训练

训练过程被实现为一个类方法，该方法会返回一个可直接使用的分词器。该分词器会先对文本中的唯一字符进行排序以构建基础词汇表，若目标词汇量小于该基础集合则不予接受，随后会像之前一样执行计数、选择和合并的循环流程。日志记录会进行限制：仅打印前20次合并的结果，之后每100次合并打印一次，同时显示字符对出现频率及当前序列长度，这样就能在不会被海量输出淹没的情况下观察压缩进度。

@classmethod
def train(
    cls,
    text,
    target_vocab_size,
    minimum_pair_frequency=2,
    verbose=True,
):
    base_tokens = sorted(
        set(text)
    )

    if target_vocab_size < len(base_tokens):
        raise ValueError(
            "target_vocab_size cannot be "
            "smaller than the number of "
            "unique characters"
        )

    working_tokens = list(text)
    vocabulary = set(base_tokens)
    merges = []

    while len(vocabulary) < target_vocab_size:
        pair_counts = Counter(
            zip(
                working_tokens[:-1],
                working_tokens[1:],
            )
        )

        if not pair_counts:
            break

        best_pair, frequency = (
            pair_counts.most_common(1)[0]
        )

        if frequency < minimum_pair_frequency:
            break

        merged_token = "".join(
            best_pair
        )

        working_tokens = merge_pair(
            tokens=working_tokens,
            pair=best_pair,
            merged_token=merged_token,
        )

        merges.append(best_pair)
        vocabulary.add(merged_token)

        if (
            verbose
            and (
                len(merges) <= 20
                or len(merges) % 100 == 0
            )
        ):
            print(
                f"Merge {len(merges):4d}: "
                f"{best_pair!r} "
                f"→ {merged_token!r}, "
                f"frequency={frequency:,}, "
                f"tokens={len(working_tokens):,}"
            )

    return cls(
        base_tokens=base_tokens,
        merges=merges,
    )

为何存在最小词对频率

如果剩余出现频率最高的词对仅出现一次，将其合并后会生成一个永远不会被重复使用的标记，这只会浪费一个词汇表位置。因此，在这种条件下循环会停止：

frequency < minimum_pair_frequency

对于这样的实验，阈值设为2就足够了：

minimum_pair_frequency = 2

在更大的语料库上，你可能需要更高的阈值，这样词汇表才会用于那些出现频率足够高、能让模型为其学习出优质嵌入的模式。

使用分词器进行编码与解码

编码过程分为三个步骤：

1. 将文本拆分为基本字符。
2. 按照学习顺序依次应用每条合并规则。
3. 查找每个生成标记的ID。

其中前半部分工作在encode_tokens函数中完成。它首先收集词汇表中缺失的字符，如果存在缺失字符则会抛出明确的错误信息，这比在处理数据集过程中遇到KeyError错误要容易调试得多。之后再依次应用合并规则。

def encode_tokens(self, text):
    unknown_characters = sorted({
        char
        for char in text
        if char not in self.stoi
    })

    if unknown_characters:
        raise ValueError(
            "Text contains characters missing "
            "from the base vocabulary: "
            f"{unknown_characters}"
        )

    tokens = list(text)

    for pair in self.merges:
        merged_token = (
            pair[0] + pair[1]
        )

        tokens = merge_pair(
            tokens=tokens,
            pair=pair,
            merged_token=merged_token,
        )

    return tokens

将生成的令牌字符串转换为ID只需进行简单的查找即可：

def encode(self, text):
    tokens = self.encode_tokens(text)

    return [
        self.stoi[token]
        for token in tokens
    ]

解码更为简单。无论是否经过合并，每个标记都存储着其所代表的确切文本，因此解码只需将它们拼接起来即可：

def decode(self, token_ids):
    return "".join(
        self.itos[int(token_id)]
        for token_id in token_ids
    )

这为BPE赋予了一个值得专门测试的特性：无损往返传输：

decode(encode(text)) = text

用一个示例句子快速测试：

sample = "The history of science."
encoded = tokenizer.encode(sample)
decoded = tokenizer.decode(encoded)
assert decoded == sample

保持和重新加载分词器

无需存储词汇表本身，因为可以根据基础标记和合并列表以确定性方式重新构建它。保存方法会将这两部分内容以及类型标签和词汇表大小以 JSON 格式写入，以便于使用。ensure_ascii=False 可确保文件中的非 ASCII 字符仍可正常显示，同时如果父目录不存在的话也会自动创建。

def save(self, path):
    data = {
        "type": "character_bpe",
        "base_tokens": self.base_tokens,
        "merges": [
            [left, right]
            for left, right in self.merges
        ],
        "vocab_size": self.vocab_size,
    }

    path = Path(path)

    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    with open(
        path,
        "w",
        encoding="utf-8",
    ) as file:
        json.dump(
            data,
            file,
            ensure_ascii=False,
            indent=2,
        )

最终会生成一个文件：

tokenizer/
└── bpe_tokenizer.json

加载时会反向执行该过程，即读取 JSON 并将两个列表一起传递给构造函数：

@classmethod
def from_file(cls, path):
    with open(
        path,
        "r",
        encoding="utf-8",
    ) as file:
        data = json.load(file)

    return cls(
        base_tokens=data[
            "base_tokens"
        ],
        merges=data["merges"],
    )

然后确认重新加载的分词器仍能对示例文本进行往返处理：

loaded_tokenizer = (
    BPETokenizer.from_file(
        "tokenizer/bpe_tokenizer.json"
    )
)

assert (
    loaded_tokenizer.decode(
        loaded_tokenizer.encode(sample)
    )
    == sample
)

在 WikiText-2 的部分数据上训练

由于简单的训练器在每次合并时都会重新扫描整个工作序列，因此使用大型词汇表对整个 WikiText-2 进行训练可能需要很长时间。可以先从可配置的训练文本前缀开始：

training_characters = 500_000

training_text = train_text[
    :training_characters
]

合理的初始词汇量应为：

target_vocab_size = 500

换言之：

Initial character vocabulary
+
learned merged tokens
=
approximately 500 tokens

当一切正常时，尝试更大的目标：

target_vocab_size = 1_000
target_vocab_size = 2_000

训练时间会随着样本规模和合并次数的增加而增长。

从每个分割中选取基础字符

如果基础词汇表仅基于训练样本构建，验证集或测试集文本中可能包含其中从未出现过的字符，从而导致这些分割的编码失败。解决方法是从三个分割中收集基础字符集，同时仍仅从训练数据中学习合并统计信息：

all_text = (
    train_text
    + val_text
    + test_text
)

base_tokens = sorted(
    set(all_text)
)

这不会将任何验证语言模式泄露到合并结果中。它仅确保数据集中的每个字符都有对应的ID。为支持这一功能，训练器会多了一个可选的base_tokens参数，当提供该参数时，它将替代训练文本中的字符使用：

@classmethod
def train(
    cls,
    text,
    target_vocab_size,
    minimum_pair_frequency=2,
    base_tokens=None,
    verbose=True,
):
    if base_tokens is None:
        base_tokens = sorted(
            set(text)
        )
    else:
        base_tokens = sorted(
            set(base_tokens)
        )

衡量获得的压缩效果

要了解BPE能带来什么好处，可比较相同文本的标记数量。以下是一个示例：

sample_text = train_text[:10_000]

其字符标记数其实就是其长度：

character_count = len(sample_text)

其BPE标记数量即为编码后ID列表的长度：

bpe_ids = tokenizer.encode(
    sample_text
)

bpe_count = len(bpe_ids)

二者之比即为压缩因子：

compression_factor = (
    character_count / bpe_count
)

结果可能如下所示：

Character tokens: 10,000
BPE tokens: 4,250
Compression: 2.35×

其压缩比为2.35，128个BPE标记所能表示的文本量大致相当于301个字符标记所能表示的量。不过压缩效果并不一致：包含常见词汇的散文压缩效果较好，而数字、生僻名称以及特殊标点符号的压缩效果则较差。

看看分词器学到了什么

这些学到的标记就是词汇表中基础字符之后的所有内容。按长度对它们进行排序并列出最长的那些标记及其编号，就能快速了解分词器认为哪些内容出现频率较高：

learned_tokens = (
    tokenizer.tokens[
        len(tokenizer.base_tokens):
    ]
)

longest_tokens = sorted(
    learned_tokens,
    key=len,
    reverse=True,
)

for token in longest_tokens[:50]:
    print(
        repr(token),
        tokenizer.stoi[token],
    )

你通常会看到类似以下的令牌：

" the"
"ing"
"tion"
" of"
"ed"
" and"

具体列表取决于训练样本、预处理方式、目标词汇表大小、最小出现频率以及平局处理规则。关于平局处理：Counter.most_common会按照首次统计的顺序返回出现频率相同的词对，因此对于给定输入，其结果是确定的；但若更改训练样本，多个频率相同的词对中哪一对获胜的结果也可能会改变。

重新编码数据集

使用已训练好的分词器，对全部三个数据集子集进行编码：

train_bpe_ids = tokenizer.encode(
    train_text
)

val_bpe_ids = tokenizer.encode(
    val_text
)

test_bpe_ids = tokenizer.encode(
    test_text
)

将 ID 列表转换为 long 类型的张量，这是嵌入层期望用于索引的类型：

train_tensor = torch.tensor(
    train_bpe_ids,
    dtype=torch.long,
)

val_tensor = torch.tensor(
    val_bpe_ids,
    dtype=torch.long,
)

test_tensor = torch.tensor(
    test_bpe_ids,
    dtype=torch.long,
)

将它们保存到各自的目录中：

data/
└── encoded_bpe/
    ├── train_ids.pt
    ├── val_ids.pt
    └── test_ids.pt

请勿修改经过字符编码的文件；保留两者可以便于后续比较，同时避免将一个模型中的标识符传递给另一个模型。

验证编码后的数据分割

首先，每个标识符的值必须介于0到词汇表大小之间（不包括词汇表大小）：

for split_name, token_ids in [
    ("train", train_tensor),
    ("validation", val_tensor),
    ("test", test_tensor),
]:
    assert token_ids.min().item() >= 0

    assert (
        token_ids.max().item()
        < tokenizer.vocab_size
    )

    print(
        split_name,
        token_ids.shape,
        token_ids.min().item(),
        token_ids.max().item(),
    )

其次，每个分割后的样本都应能往返还原为原始文本：

for text in [
    train_text[:1_000],
    val_text[:1_000],
    test_text[:1_000],
]:
    encoded = tokenizer.encode(text)
    decoded = tokenizer.decode(encoded)

    assert decoded == text

将BPE集成到模型中

模型架构无需任何结构上的修改。仅需调整配置中的词汇表大小：

config = MiniGPTConfig(
    vocab_size=bpe_tokenizer.vocab_size,
    block_size=128,
    embedding_dim=128,
    num_heads=4,
    num_layers=4,
    expansion_factor=4,
    dropout=0.1,
)

将数据加载指向新文件：

train_ids = torch.load(
    "data/encoded_bpe/train_ids.pt"
)

从现在开始，该模型将预测BPE标记而非字符。

旧的检查点将无法加载

字符模型通常的词汇量范围为：

100 to 300 character tokens

BPE模型通常具有类似以下结构：

500 to 2,000 subword tokens

令牌嵌入和输出头都依赖于词汇表大小，因此它们的结构会发生变化。字符检查点的结构为：

token embedding: [character_vocab_size, C]

而BPE模型期望的是：

token embedding: [bpe_vocab_size, C]

此处 C 表示嵌入维度。由于形状不匹配，旧权重无法直接加载。即便可以加载，不同分词器下的 ID 含义也各不相同。你需要训练一个全新的模型，此外将检查点分开保存也会更有帮助：

checkpoints_char/
checkpoints_bpe/

或者，等价地：

checkpoints/char/
checkpoints/bpe/

标记化器的指纹识别

形状不匹配会产生明显差异；而将检查点与大小相同但不同的BPE标记化器配对则不会产生明显变化，因为每个ID都只会对应到错误的文本。该指纹技术会对基础标记的紧凑JSON格式进行哈希处理，并与SHA-256合并，因此任何变化都会生成新的值。

import hashlib


def tokenizer_fingerprint(tokenizer):
    payload = json.dumps(
        {
            "base_tokens": (
                tokenizer.base_tokens
            ),
            "merges": [
                list(pair)
                for pair in tokenizer.merges
            ],
        },
        ensure_ascii=False,
        separators=(",", ":"),
    ).encode("utf-8")

    return hashlib.sha256(
        payload
    ).hexdigest()

将指纹存储在每个检查点中：

checkpoint[
    "tokenizer_fingerprint"
] = tokenizer_fingerprint(
    tokenizer
)

在加载时，计算当前分词器的指纹，如果其与检查点中保存的指纹不一致，则拒绝继续处理。

你需要了解的性能限制

编码器会遍历整个分词序列，重新应用每一条合并规则：

for pair in self.merges:
    tokens = merge_pair(...)

每次合并都会对文本进行一次完整遍历，因此包含数千次合并的大型文本编码速度会很慢。该实现刻意优先考虑清晰性、正确性及可检查性，而非速度。生产环境中的分词器会预先将文本拆分为单词，缓存重复的分割结果，并按优先级进行合并；而在学习场景中，透明版本更具优势。

完整程序

将所有内容整合到一个脚本中：

build_bpe_tokenizer.py

下面的完整程序整合了配对辅助函数、分词器类、验证辅助函数以及命令行入口点。与示例代码相比有两个不同之处：保存时现在会存储指纹，并通过临时文件进行写入，之后再将临时文件重命名为目标文件，这样即使运行过程中中断，也不会留下未写完的文件；此外，训练过程还会验证基础词汇表是否涵盖了训练文本。主函数会在配置好的前缀上进行训练，保存并重新加载分词器，检查指纹，对每个分割数据集进行往返处理并保存，验证编号范围，最后输出统计信息。

“最长已学习标记”循环中的一行似乎出现了错误：f"{loaded_tokenizer.stoi5d} "会引发AttributeError异常，因为并不存在stoi5d这个函数。正确的表达式很可能是f"{loaded_tokenizer.stoi[token]:5d} "，它能将每个标记的ID右补为五位数字并输出。在运行脚本之前请先修正那一行。

import argparse
import hashlib
import json
from collections import Counter
from pathlib import Path

import torch


# ==================================================
# Pair Operations
# ==================================================

def count_pairs(tokens):
    return Counter(
        zip(tokens[:-1], tokens[1:])
    )


def merge_pair(
    tokens,
    pair,
    merged_token,
):
    output = []
    index = 0

    while index < len(tokens):
        if (
            index < len(tokens) - 1
            and tokens[index] == pair[0]
            and tokens[index + 1] == pair[1]
        ):
            output.append(merged_token)
            index += 2
        else:
            output.append(tokens[index])
            index += 1

    return output


# ==================================================
# BPE Tokenizer
# ==================================================

class BPETokenizer:
    def __init__(
        self,
        base_tokens,
        merges,
    ):
        self.base_tokens = list(
            base_tokens
        )

        self.merges = [
            tuple(pair)
            for pair in merges
        ]

        self.tokens = list(
            self.base_tokens
        )

        for left, right in self.merges:
            merged_token = left + right

            if merged_token not in self.tokens:
                self.tokens.append(
                    merged_token
                )

        self.stoi = {
            token: index
            for index, token
            in enumerate(self.tokens)
        }

        self.itos = {
            index: token
            for token, index
            in self.stoi.items()
        }

        self.vocab_size = len(
            self.tokens
        )

    @classmethod
    def train(
        cls,
        text,
        target_vocab_size,
        minimum_pair_frequency=2,
        base_tokens=None,
        verbose=True,
    ):
        if not text:
            raise ValueError(
                "Training text cannot be empty"
            )

        if base_tokens is None:
            base_tokens = sorted(
                set(text)
            )
        else:
            base_tokens = sorted(
                set(base_tokens)
            )

        missing_characters = (
            set(text) - set(base_tokens)
        )

        if missing_characters:
            raise ValueError(
                "Base vocabulary is missing "
                "training characters"
            )

        if (
            target_vocab_size
            < len(base_tokens)
        ):
            raise ValueError(
                "target_vocab_size cannot be "
                "smaller than the number of "
                "base tokens"
            )

        working_tokens = list(text)
        vocabulary = set(base_tokens)
        merges = []

        print(
            "Initial characters:",
            len(working_tokens),
        )

        print(
            "Base vocabulary size:",
            len(base_tokens),
        )

        while (
            len(vocabulary)
            < target_vocab_size
        ):
            pair_counts = count_pairs(
                working_tokens
            )

            if not pair_counts:
                print(
                    "No adjacent pairs remain."
                )
                break

            best_pair, frequency = (
                pair_counts.most_common(1)[0]
            )

            if (
                frequency
                < minimum_pair_frequency
            ):
                print(
                    "Stopping because the most "
                    "frequent pair occurs only "
                    f"{frequency} time(s)."
                )
                break

            merged_token = "".join(
                best_pair
            )

            working_tokens = merge_pair(
                tokens=working_tokens,
                pair=best_pair,
                merged_token=merged_token,
            )

            merges.append(best_pair)
            vocabulary.add(merged_token)

            merge_number = len(merges)

            if (
                verbose
                and (
                    merge_number <= 20
                    or merge_number % 100 == 0
                )
            ):
                print(
                    f"Merge {merge_number:4d}: "
                    f"{best_pair!r} "
                    f"→ {merged_token!r}, "
                    f"frequency={frequency:,}, "
                    f"sequence length="
                    f"{len(working_tokens):,}, "
                    f"vocabulary="
                    f"{len(vocabulary):,}"
                )

        tokenizer = cls(
            base_tokens=base_tokens,
            merges=merges,
        )

        print(
            "Final vocabulary size:",
            tokenizer.vocab_size,
        )

        print(
            "Final training sequence length:",
            len(working_tokens),
        )

        return tokenizer

    @classmethod
    def from_file(cls, path):
        with open(
            path,
            "r",
            encoding="utf-8",
        ) as file:
            data = json.load(file)

        return cls(
            base_tokens=data[
                "base_tokens"
            ],
            merges=data["merges"],
        )

    def encode_tokens(self, text):
        unknown_characters = sorted({
            char
            for char in text
            if char not in self.stoi
        })

        if unknown_characters:
            raise ValueError(
                "Text contains characters "
                "missing from the base vocabulary: "
                f"{unknown_characters}"
            )

        tokens = list(text)

        for left, right in self.merges:
            tokens = merge_pair(
                tokens=tokens,
                pair=(left, right),
                merged_token=left + right,
            )

        return tokens

    def encode(self, text):
        return [
            self.stoi[token]
            for token in self.encode_tokens(
                text
            )
        ]

    def decode(self, token_ids):
        return "".join(
            self.itos[int(token_id)]
            for token_id in token_ids
        )

    def save(self, path):
        path = Path(path)

        path.parent.mkdir(
            parents=True,
            exist_ok=True,
        )

        data = {
            "type": "character_bpe",
            "base_tokens": (
                self.base_tokens
            ),
            "merges": [
                [left, right]
                for left, right
                in self.merges
            ],
            "vocab_size": (
                self.vocab_size
            ),
            "fingerprint": (
                self.fingerprint()
            ),
        }

        temporary_path = path.with_suffix(
            ".tmp"
        )

        with open(
            temporary_path,
            "w",
            encoding="utf-8",
        ) as file:
            json.dump(
                data,
                file,
                ensure_ascii=False,
                indent=2,
            )

        temporary_path.replace(path)

    def fingerprint(self):
        payload = json.dumps(
            {
                "base_tokens": (
                    self.base_tokens
                ),
                "merges": [
                    list(pair)
                    for pair in self.merges
                ],
            },
            ensure_ascii=False,
            separators=(",", ":"),
        ).encode("utf-8")

        return hashlib.sha256(
            payload
        ).hexdigest()


# ==================================================
# Validation
# ==================================================

def validate_round_trip(
    tokenizer,
    text,
    split_name,
):
    encoded = tokenizer.encode(text)
    decoded = tokenizer.decode(encoded)

    if decoded != text:
        raise AssertionError(
            f"Round-trip failed for "
            f"{split_name}"
        )

    print(
        f"{split_name} round-trip passed: "
        f"{len(text):,} characters "
        f"→ {len(encoded):,} tokens"
    )

    return encoded


def save_token_tensor(
    token_ids,
    path,
):
    path = Path(path)

    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    tensor = torch.tensor(
        token_ids,
        dtype=torch.long,
    )

    temporary_path = path.with_suffix(
        ".tmp"
    )

    torch.save(
        tensor,
        temporary_path,
    )

    temporary_path.replace(path)

    return tensor


# ==================================================
# Arguments
# ==================================================

def parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "Train a simple educational BPE "
            "tokenizer on WikiText-2"
        )
    )

    parser.add_argument(
        "--processed-dir",
        type=str,
        default="data/processed",
    )

    parser.add_argument(
        "--tokenizer-output",
        type=str,
        default=(
            "tokenizer/bpe_tokenizer.json"
        ),
    )

    parser.add_argument(
        "--encoded-output-dir",
        type=str,
        default="data/encoded_bpe",
    )

    parser.add_argument(
        "--target-vocab-size",
        type=int,
        default=500,
    )

    parser.add_argument(
        "--training-characters",
        type=int,
        default=500_000,
    )

    parser.add_argument(
        "--minimum-pair-frequency",
        type=int,
        default=2,
    )

    return parser.parse_args()


# ==================================================
# Main Program
# ==================================================

def main():
    args = parse_args()

    processed_dir = Path(
        args.processed_dir
    )

    train_path = (
        processed_dir
        / "wikitext2_train_clean.txt"
    )

    val_path = (
        processed_dir
        / "wikitext2_val_clean.txt"
    )

    test_path = (
        processed_dir
        / "wikitext2_test_clean.txt"
    )

    for path in [
        train_path,
        val_path,
        test_path,
    ]:
        if not path.exists():
            raise FileNotFoundError(
                f"Required file not found: "
                f"{path}"
            )

    train_text = train_path.read_text(
        encoding="utf-8"
    )

    val_text = val_path.read_text(
        encoding="utf-8"
    )

    test_text = test_path.read_text(
        encoding="utf-8"
    )

    print(
        "Training characters:",
        len(train_text),
    )

    print(
        "Validation characters:",
        len(val_text),
    )

    print(
        "Test characters:",
        len(test_text),
    )

    all_characters = sorted(
        set(
            train_text
            + val_text
            + test_text
        )
    )

    training_character_count = min(
        args.training_characters,
        len(train_text),
    )

    bpe_training_text = train_text[
        :training_character_count
    ]

    print(
        "\nBPE training characters:",
        len(bpe_training_text),
    )

    tokenizer = BPETokenizer.train(
        text=bpe_training_text,
        target_vocab_size=(
            args.target_vocab_size
        ),
        minimum_pair_frequency=(
            args.minimum_pair_frequency
        ),
        base_tokens=all_characters,
        verbose=True,
    )

    tokenizer.save(
        args.tokenizer_output
    )

    print(
        "\nTokenizer saved to:",
        args.tokenizer_output,
    )

    print(
        "Tokenizer fingerprint:",
        tokenizer.fingerprint(),
    )

    loaded_tokenizer = (
        BPETokenizer.from_file(
            args.tokenizer_output
        )
    )

    assert (
        loaded_tokenizer.fingerprint()
        == tokenizer.fingerprint()
    )

    print(
        "Tokenizer reload test passed."
    )

    print(
        "\nEncoding dataset splits..."
    )

    train_ids = validate_round_trip(
        tokenizer=loaded_tokenizer,
        text=train_text,
        split_name="Training",
    )

    val_ids = validate_round_trip(
        tokenizer=loaded_tokenizer,
        text=val_text,
        split_name="Validation",
    )

    test_ids = validate_round_trip(
        tokenizer=loaded_tokenizer,
        text=test_text,
        split_name="Test",
    )

    output_dir = Path(
        args.encoded_output_dir
    )

    train_tensor = save_token_tensor(
        train_ids,
        output_dir / "train_ids.pt",
    )

    val_tensor = save_token_tensor(
        val_ids,
        output_dir / "val_ids.pt",
    )

    test_tensor = save_token_tensor(
        test_ids,
        output_dir / "test_ids.pt",
    )

    for split_name, tensor in [
        ("Training", train_tensor),
        ("Validation", val_tensor),
        ("Test", test_tensor),
    ]:
        assert tensor.min().item() >= 0

        assert (
            tensor.max().item()
            < loaded_tokenizer.vocab_size
        )

        print(
            f"{split_name} tensor: "
            f"shape={tuple(tensor.shape)}, "
            f"min={tensor.min().item()}, "
            f"max={tensor.max().item()}"
        )

    character_tokens = len(
        train_text
    )

    bpe_tokens = len(
        train_ids
    )

    compression_factor = (
        character_tokens / bpe_tokens
    )

    reduction_percentage = (
        1.0
        - bpe_tokens / character_tokens
    ) * 100.0

    print(
        "\nCharacter token count:",
        f"{character_tokens:,}",
    )

    print(
        "BPE token count:",
        f"{bpe_tokens:,}",
    )

    print(
        "Compression factor:",
        f"{compression_factor:.2f}x",
    )

    print(
        "Token-count reduction:",
        f"{reduction_percentage:.2f}%",
    )

    learned_tokens = (
        loaded_tokenizer.tokens[
            len(
                loaded_tokenizer.base_tokens
            ):
        ]
    )

    longest_tokens = sorted(
        learned_tokens,
        key=len,
        reverse=True,
    )

    print("\nLongest learned tokens:")

    for token in longest_tokens[:30]:
        print(
            f"{loaded_tokenizer.stoi5d} "
            f"{token!r}"
        )

    sample = (
        "Artificial intelligence "
        "and machine learning"
    )

    unknown = [
        char
        for char in sample
        if char not in (
            loaded_tokenizer.stoi
        )
    ]

    if not unknown:
        sample_tokens = (
            loaded_tokenizer.encode_tokens(
                sample
            )
        )

        sample_ids = (
            loaded_tokenizer.encode(
                sample
            )
        )

        print("\nSample text:")
        print(sample)

        print("\nSample BPE tokens:")
        print(sample_tokens)

        print("\nSample token IDs:")
        print(sample_ids)

        print("\nDecoded sample:")
        print(
            loaded_tokenizer.decode(
                sample_ids
            )
        )

    print(
        "\nBPE tokenizer pipeline "
        "completed successfully."
    )


if __name__ == "__main__":
    main()

运行它

默认配置为一次小型、用于教学的运行：

python build_bpe_tokenizer.py

要学习更多词汇：

python build_bpe_tokenizer.py \
  --target-vocab-size 1000

若要使用更多文本进行训练：

python build_bpe_tokenizer.py \
  --training-characters 1000000 \
  --target-vocab-size 1000

如需快速进行仅用于确认流水线是否正常的测试：

python build_bpe_tokenizer.py \
  --training-characters 100000 \
  --target-vocab-size 300

如果您的shell不支持反斜杠换行，那么请将每个命令写在单独的一行上。

输出效果

执行后可能会显示类似这样的内容：

Training characters: 10,000,000
Validation characters: 1,100,000
Test characters: 1,200,000

BPE training characters: 500,000
Initial characters: 500,000
Base vocabulary size: 200

Merge    1: (' ', 't') → ' t', frequency=19,842
Merge    2: ('h', 'e') → 'he', frequency=17,415
Merge    3: (' t', 'he') → ' the', frequency=14,108
...
Final vocabulary size: 500

最初的合并过程完全符合人们对英语的预期：一个空格，接着是 t，然后是 he，最后这两个字符组合成 " the"。编码后的结果大致如下：

Character token count: 10,000,000
BPE token count: 4,500,000
Compression factor: 2.22x
Token-count reduction: 55.00%

具体数值会有所不同。请注意，使用该编码器对完整训练数据集进行往返处理会在词汇量较大时占据大量运行时间。

模型及其上下文窗口有何变化

真实句子中的标记数量

以这个句子为例：

The history of artificial intelligence

字符分词器大约会生成这么多个词元：

38 tokens

BPE分词器可能会生成如下这样的整词标记：

["The", " history", " of", " artificial", " intelligence"]

或者根据它所学的合并规则，它可能会生成更小的片段。无论哪种情况，数量都会落在类似以下的范围内：

5 to 15 tokens

这种分词并非人工设计，而是源自训练语料库的统计结果。

每个上下文窗口容纳更多文本

假设模型使用这样的块大小：

block_size = 128

使用字符表示时，128个标记大约对应128个字符。而采用BPE后，同样的128个标记就能覆盖数百个字符。模型能够在不增加注意力计算量的情况下获取更多语言上下文。这一点很重要，因为自注意力的计算成本大致与以下因素相关：

T²

此处 T 表示序列长度，因此将 T 翻倍大约会使注意力计算成本增加四倍，而更好的分词方式则能在相同的 T 下容纳更多文本。

词汇表大小之间的权衡

更大的 BPE 词汇表通常意味着更短的序列，但也会使某些矩阵的规模变大。词元嵌入表需要存储这么多的参数：

vocab_size × embedding_dim

在权重绑定的情况下，输出头会重用同一个矩阵，而不会添加第二个矩阵。以较小的词汇表为例：

Vocabulary:          500
Embedding dimension: 128
Embedding parameters: 64,000

使用更大的那个：

Vocabulary: 8,000
Embedding dimension: 128
Embedding parameters: 1,024,000

因此，更大的词汇表虽然能缩短序列长度，但会增加嵌入参数、输出对数概率的成本、内存占用以及分词器训练时间。每个额外的标记也需要足够的出现次数才能学习到良好的嵌入表示，所以在像 WikiText-2 这样规模的语料库中，过大的词汇表会导致许多标记几乎无法被充分训练。合适的词汇表大小取决于数据集和模型。

使用 BPE 标记重新训练模型

需更新训练脚本以加载分词器文件：

tokenizer/bpe_tokenizer.json

编码后的分割结果：

data/encoded_bpe/train_ids.pt
data/encoded_bpe/val_ids.pt
data/encoded_bpe/test_ids.pt

然后根据分词器的词汇表大小来构建配置：

config = MiniGPTConfig(
    vocab_size=bpe_tokenizer.vocab_size,
    block_size=128,
    embedding_dim=128,
    num_heads=4,
    num_layers=4,
    expansion_factor=4,
    dropout=0.1,
)

并将检查点写入它们自己的目录中：

checkpoints/bpe/

除了这些变化之外，训练循环基本上可以保持原样。

公平比较两种模型

为了进行有意义的比较，尽可能固定架构（宽度、深度、输出头数量）、数据、优化器、调度策略、训练预算以及评估方法。即便如此，按相同的训练步数进行比较也会产生误导，因为一批BPE标记所包含的文本量远多于一批字符。更合适的比较维度是：

Total characters processed
Total bytes processed
Training time
Compute budget

困惑度是按预测单元计算的，因此字符困惑度和BPE困惑度无法直接比较。若要获得可比较的数值，需将两者都转换为每字符比特数；否则应通过生成质量、计算成本以及适合分词器的指标来评估各模型。

常见错误

为每个分割单独构建分词器

这是错误的：

Train tokenizer from training text
Validation tokenizer from validation text
Test tokenizer from test text

训练数据、验证数据以及测试数据都必须使用同一个分词器进行编码；否则，在不同的数据划分中，相同的ID可能会代表不同的含义。

错误顺序的合并操作

合并操作是相互依赖的，因此必须按照它们被学习到的确切顺序来执行：

for pair in self.merges:
    apply(pair)

对规则进行排序，例如按字母顺序以便于阅读，会改变文本的分割方式。

覆盖字符编码文件

将BPE ID保存在单独的目录中：

data/encoded_bpe/

并将字符编号保持原位：

data/encoded/

从验证文本中学习合并规则

合并统计数据应仅来自训练文本。为确保覆盖范围，从每个分割部分中提取基础字符集是可行的，但合并频率必须仅反映训练语料库的情况，否则验证数据将会偏高。

期望学到的标记为完整单词

BPE会轻松学习诸如以下的片段：

"ing"
"tion"
" th"
" arti"

子词标记遵循出现频率而非人类设定的词界，这是预期中的行为。

认为词汇量越大越好是错误的

更大的词汇量会增加参数开销，还会产生出现频率低、训练效果差的标记。建议尝试多种大小并进行测试。同样，切勿直接比较不同分词器的原始困惑度。

练习

1. 比较三种不同的词汇量。使用以下三种大小进行训练，然后比较训练时间、序列长度、压缩率、最长标记的长度以及嵌入向量的维度。

300 tokens
500 tokens
1,000 tokens

1. 研究早期的合并规则。打印前50条合并规则。其中是否存在空白字符模式？常见的后缀？标点符号组合？有哪些合并规则让你感到意外？
2. 用两种方式对一个句子进行分词。使用字符分词器和BPE分词器对以下句子进行编码，然后比较生成的词元、标识符、词元数量以及解码结果。

Artificial intelligence is transforming technology.

1. 测试生僻词。对以下单词进行编码，观察BPE是如何用已知片段组合出陌生单词的。

electromagnetism
neurocomputational
uncharacteristically

1. 衡量上下文覆盖范围。解码每个分词器处理后的前128个标记，查看每个窗口覆盖的文本量。
2. 训练BPE模型。使用下方目录中的数据进训练，并在相同的计算资源下将其生成的样本与字符模型生成的样本进行比较。

data/encoded_bpe/

1. 强制使用指纹识别机制。在每个检查点中存储分词器的指纹，当检查点与当前分词器不一致时拒绝进行推理。
2. 加快编码速度。用更高效的编码方式替代逐次合并的编码器，可利用合并等级、堆结构、链式词表、单词缓存或尝试法等。

值得保留的分词规则

• 子词介于灵活但长度较长的字符与紧凑但易出问题的单词之间。
• BPE分词器完全由其基础词元及按顺序的合并规则定义；应精确保存这些内容并生成指纹。
• 编码时按顺序重新执行合并操作，解码时则将结果直接拼接，往返处理过程无信息损失。
• 仅使用一个分词器处理所有分词任务，合并规则完全从训练数据中学习得到。

新的分词器意味着新的词元文件、更大的词汇量以及重新训练过的模型。

词汇量的增加会以更短的序列为代价，换取更大的嵌入层和输出层，而困惑度不会在不同分词器之间传递。

从原始文本到重新训练后的模型的完整流程如下：

WikiText training text
→ character vocabulary
→ pair frequencies
→ repeated merge operations
→ BPE vocabulary and merge rules
→ encode all splits
→ save token tensors
→ train a new Mini-GPT

更关键的是，分词并非简单的预处理步骤。它决定了模型要学习的单元是什么、多少文本能纳入其上下文窗口，以及其输出预测问题的规模大小。一旦小型GPT开始使用BPE分词，接下来的自然步骤就是探究同样的分词机制如何发展成现代生成式AI系统：

Next-token predictor
→ pretrained language model
→ instruction-following model
→ retrieval-augmented application
→ tool-using GenAI system

相关阅读

• 在 PyTorch 中调试小型 GPT：能够隔离每处故障的测试方法：详细介绍在 PyTorch 中逐步调试字符级 GPT 的流程，涵盖从令牌 ID 和目标偏移，到梯度计算、NaN 损失以及检查点的处理。
• 从零构建 AI 智能体：Patterns、ReAct 与 LangGraph：了解 AI 智能体的核心概念，包括规划、工具使用、反思机制以及 ReAct 模式，同时探讨 LangChain 和 LangGraph 如何帮助手动构建此类智能体。

从GPT-1到推理模型：每一代对开发者带来了哪些变化：追溯GPT系列从2018年的预训练阶段到推理模型的发展历程，了解每一代新增了哪些技术理念，同时学习如何调用GPT-4o的视觉功能以及Python的结构化输出功能。

在8 GB内存的笔记本电脑上从零开始训练小型大语言模型：介绍如何在配置较为有限的硬件上搭建用于训练Transformer模型的数据处理流程，包括BPE分词器、TinyStories数据集、基于磁盘的令牌存储方式以及下一个令牌的批量处理方法。