2 Byte-Pair Encoding Tokenizer#

在第一部分，我们将训练并实现一个字节级字节对编码分词器（byte-level BPE tokenizer）。具体来说，我们将任意（Unicode）字符串表示为字节序列，并基于该字节序列训练我们的 BPE tokenizer。之后，我们将使用该 tokenizer 将 text（字符串）编码为 tokens（整数序列），以用于语言建模。

2.1 The Unicode Standard#

请自行详细了解 python 中的 ord() chr()函数的作用及其背景。

2.2 Unicode Encodings#

在本文中，我们将广泛使用UTF-8编码

将字符串编码成为bytes，我们可以查看以下示例：

>>> test_string = "hello! こんにちは!"
>>> utf8_encoded = test_string.encode("utf-8")
>>> print(utf8_encoded)
b'hello! \xe3\x81\x93\xe3\x82\x93\xe3\x81\xab\xe3\x81\xa1\xe3\x81\xaf!'
>>> print(type(utf8_encoded))
<class 'bytes'>
>>> # Get the byte values for the encoded string (integers from 0 to 255).
>>> list(utf8_encoded)
[104, 101, 108, 108, 111, 33, 32, 227, 129, 147, 227, 130, 147, 227, 129, 171, 227, 129,
161, 227, 129, 175, 33]
>>> # One byte does not necessarily correspond to one Unicode character!
>>> print(len(test_string))
13
>>> print(len(utf8_encoded))
23
>>> print(utf8_encoded.decode("utf-8"))
hello! こんにちは!

2.3 Subword Tokenization#

虽然 byte-level tokenization 可以缓解 word-level tokenizers 面临的词汇表外问题，但将文本分词为字节会导致输入序列极长。

这会减慢模型训练的速度，因为在 word-level language model 中，一个包含10个单词的句子可能只有10个词，但在字符级模型中（取决于单词的长度），这个句子可能包含50个甚至更多的词。处理这些更长的序列需要在模型的每一步进行更多的计算。此外，对字节序列进行语言建模也很困难，因为更长的输入序列会在数据中产生长期依赖关系。

Subword tokenization 是 word-level tokenizers 和 byte-level tokenizers 之间的折中方案。请注意，byte-level tokenizers的词汇表有256个条目（字节值为0到225）。Subword tokenizer 以更大的词汇表为代价，换取对输入字节序列更好的压缩效果。

例如，如果字节序列b'the'在我们的原始文本训练数据中经常出现，那么在词汇表中为其分配一个条目，就会将这个3个词的序列简化为一个单独的词。

我们如何选择这些 subword unit 来扩充我们的词汇表呢？Sennrich 等人建议使用字节对编码（BPE；Gage，1994），这是一种压缩算法，它通过迭代替换（“合并”）最频繁出现的字节对，用一个全新的、未使用的索引来表示。请注意，该算法将子词标记添加到我们的词汇表中，以最大限度地压缩输入序列——如果一个词在我们的输入文本中出现的次数足够多，它将被表示为一个单一的子词单元。

通过BPE构建词汇表的 subword tokenizers 通常被称为 BPE tokenizers 。在本作业中，我们将实现一个byte-level BPE tokenizer，其中词汇表项是字节或合并的字节序列，这既有利于处理词汇表外词汇，也有利于管理输入序列长度。构建 BPE tokenizers vocabulary 的过程被称为“训练”BPE分词器。

2.4 BPE Tokenizer Training#

BPE分词器的训练过程包括三个主要步骤

1 Vocabulary initialization#

Tokenizer vocabulary 是从字节串标记到整数ID的一一映射。由于我们正在训练一个 byte-level BPE tokenizer，因此我们的初始 vocabulary 只是所有字节的集合。由于有256种可能的字节值，因此我们的初始词汇表大小为256。

2 Pre-tokenization#

一旦你有了词汇表，原则上，你可以统计文本中相邻字节出现的频率，并从最频繁的字节对开始合并它们。然而，这在计算上相当昂贵，因为每次合并时，我们都必须对语料库进行一次完整的遍历。

此外，直接在整个语料库中合并字节可能会导致标记之间仅在标点符号上有所不同（例如，“dog!”与“dog.”）。尽管这些标记在语义上可能高度相似（因为它们仅在标点符号上有所不同），但它们会获得完全不同的 token IDs。

为了避免这种情况，我们对语料库进行了预分词。你可以将此视为对语料库进行的一种粗粒度分词(coarse-grained tokenization)，这有助于我们统计字符对出现的频率。

例如，单词“text”可能是一个预分词，出现了10次。在这种情况下，当我们统计字符“t”和“e”相邻出现的频率时，我们会发现单词“text”中“t”和“e”是相邻的，因此我们可以将它们的计数增加10，而无需遍历整个语料库。由于我们正在训练一个 byte-level BPE model，因此每个预分词都表示为一个UTF-8字节序列。

Sennrich等人最初的 BPE 实现通过简单地在空白处进行分割（即s.split(” “））来进行预分词。相比之下，我们将使用来自 github.com/openai/tiktoken/pull/234/files 的基于正则表达式的预分词器（GPT-2所使用；Radford等人，2019）：

>>> PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

使用此预分词器对一些文本进行交互式分词，以便更好地了解其行为，这可能很有用：

>>> # requires `regex` package
>>> import regex as re
>>> re.findall(PAT, "some text that i'll pre-tokenize")

output:

['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

3 Compute BPE merges#

现在我们已经将输入文本转换为 pre-tokens，并将每个 pre-tokens 表示为UTF-8字节序列，接下来我们可以计算BPE合并（即训练BPE分词器）。

从宏观上看，BPE算法会迭代地统计每一对字节，并找出出现频率最高的字节对（“A”，“B”）。然后，每次出现这个最频繁的字节对（“A”，“B”）时，都会将其合并，即替换为新的 token “AB”。这个新合并的 token 将被添加到我们的词汇表中；因此，BPE训练后的最终词汇表大小为初始词汇表（在我们的例子中为256个）加上训练过程中执行的BPE合并操作次数。

为了提高BPE训练的效率，我们不考虑跨越 pre-token 边界的字节对。在计算合并时，通过优先选择字典序更大的字节对，确定性地打破字节对频率的平局。例如，如果字节对（“A”，“B”）、（“A”，“C”）、（“B”，“ZZ”）和（“BA”，“A”）都具有最高频率，我们会合并（“BA”，“A”）：

>>> max([("A", "B"), ("A", "C"), ("B", "ZZ"), ("BA", "A")])

output:

('BA', 'A')

4 Special tokens#

通常，一些字符串（例如，<|endoftext|>）用于编码元数据(metadata)（例如，文档之间的边界）。在编码文本时，通常希望将某些字符串视为**“special tokens”**，这些 tokens 永远不应被拆分为多个token（即，始终作为 single token保留）。

例如，序列结束字符串<|endoftext|>应始终作为 single token（即，single integer ID）保留，以便我们知道何时停止从语言模型生成。这些特殊标记必须添加到词汇表中，以便它们具有相应的固定 token ID。

Sennrich等人的算法1包含了一个BPE分词器训练的低效实现（基本上遵循了我们上面概述的步骤）。作为第一个练习，实现并测试这个函数来检验你的理解程度，可能会很有帮助。

Example(bpe_example): BPE training example#

这是Sennrich等人给出的一个程式化示例。考虑一个语料库，其中包含以下文本：

low low low low low
lower lower widest widest widest
newest newest newest newest newest newest

词汇里面还有一个 special token: <|endoftext|>

Vocabulary#

我们使用 special token <|endoftext|>和256个字节值来初始化词汇表。

Pre-tokenization#

为了简化并专注于合并过程，在本例中我们假设 pretokenization 只是简单地根据空格进行分割。当我们进行预分词和计数时，最终会得到一个频率表(frequency table)。

{low: 5, lower: 2, widest: 3, newest: 6}

将其表示为dict[tuple[bytes], int]（例如{(l,o,w): 5 …}）会很方便。请注意，在Python中，即使单个字节也是一个 bytes 对象。Python中没有表示单个字节的 byte 类型，就像Python中没有表示单个字符的 char 类型一样。

Merges#

我们首先查看每一对连续的字节，并计算它们出现的单词的频率总和，结果为

{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 9, st: 9, ne: 6, ew: 6}

其中，('es')和('st')出现的频率相同，因此我们选择字典序更大的组合，即('st')。

然后，我们将预处理后的词元进行合并，最终得到

{(l,o,w): 5, (l,o,w,e,r): 2, (w,i,d,e,st): 3, (n,e,w,e,st): 6}

在第二轮中，我们发现(e, st)是最常见的组合（出现9次），我们会将其合并为

{(l,o,w): 5, (l,o,w,e,r): 2, (w,i,d,est): 3, (n,e,w,est): 6}

以此类推，我们最终得到的合并序列将是

['s t', 'e st', 'o w', 'l ow', 'w est', 'n e', 'ne west', 'w i', 'wi d', 'wid est', 'low e', 'lowe r']

如果我们进行6次合并，我们得到

['s t', 'e st', 'o w', 'l ow', 'w est', 'n e']

我们的词汇表元素将为

[<|endoftext|>, [...256 BYTE CHARS], st, est, ow, low, west, ne]

使用这个词汇表和合并集，单词 newest 将被标记为[ne, west]。

2.5 Experimenting with BPE Tokenizer Training#

让我们在 TinyStories数据集上训练一个 byte-level BPE tokenizer。有关查找/下载该数据集的说明，请参阅第1节。在开始之前，我们建议您先浏览一下 TinyStories 数据集，以便对数据内容有所了解。

1 Parallelizing pre-tokenization#

你会发现，预分词步骤是一个主要瓶颈。你可以通过使用内置库multiprocessing 对代码进行并行化来加速预分词。具体来说，我们建议在预分词的并行实现中，对语料库进行分块，同时确保分块边界出现在special token 的开头。你可以直接使用以下链接中的示例代码来获取分块边界，然后利用这些边界在进程间分配工作：

https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/pretokenization_example.py

这种分块方式始终有效，因为我们从不希望跨文档边界进行合并。在作业中，你可以始终采用这种方式进行拆分。不用担心遇到接收到的语料库非常大且不包含<|endoftext|>的极端情况。

2 Pre-tokenization之前去掉special tokens#

在使用正则表达式模式（使用re.finditer）进行预分词之前，您应该从语料库（或使用并行实现时的语块）中移除所有 special tokens。确保在 special tokens 处进行分割，以避免它们分隔的文本发生合并。

例如，如果您的语料库（或语块）类似于[Doc 1]<|endoftext|>[Doc 2]，则应在特殊标记<|endoftext|>处进行分割，并对[Doc 1]和[Doc 2]分别进行预分词，以避免文档边界发生合并。

这可以通过使用re.split函数，并将“|”.join(special_tokens)作为分隔符来实现（由于特殊标记中可能包含|，因此要谨慎使用re.escape）。测试脚本test_train_bpe_special_tokens将对此进行测试。

3 优化合并步骤#

上述程式化示例中BPE训练的原始实现速度较慢，因为每次合并时，它都需要遍历所有字节对以找出出现最频繁的字节对。然而，每次合并后，只有与合并后的字节对重叠的字节对的计数才会发生变化。因此，可以通过为所有字节对建立索引并增量更新这些计数来提高BPE训练的速度，而不是显式地遍历每一对字节来计算字节对的频率。使用这种缓存过程可以显著提高速度，但请注意，在Python中，BPE训练的合并部分是不可并行化的。

Problem (train_bpe): BPE Tokenizer Training#

编写一个函数，该函数接受一个输入文本文件的路径，用于训练一个（字节级）BPE分词器。

Input#

input_path: str

指向包含 BPE tokenizer 训练数据的文本文件的路径。

vocab_size: int

一个正整数，定义最终词汇表的最大大小（包括初始字节词汇表、合并产生的词汇表项以及任何 special tokens）。

special_tokens: list[str]

这是一个要添加到词汇表中的字符串列表。这些 special tokens 不会影响BPE（字节级预训练编码）训练。

Output#

vocab: dict[int, bytes]

分词器词汇表，即一个从整数（词汇表中的词元ID）到字节（词元字节）的映射。

merges: list[tuple[bytes, bytes]]

训练过程中产生的 BPE 合并列表。每个列表项是一个 tuple of bytes（,），表示与已合并。合并应按创建顺序排序。

备注#

为了使用我们提供的测试来测试你的BPE训练功能，你首先需要在[adapters.run_train_bpe]处实现测试适配器。

然后，运行uv run pytest tests/test_train_bpe.py，你的实现应该能够通过所有测试。

作为可选方案（这可能需要大量时间投入），你可以使用一些系统语言（例如C++（可以考虑使用cppyy）或Rust（使用PyO3））来实现训练方法的关键部分。

如果你这样做，请注意哪些操作需要复制数据，哪些操作需要直接从Python内存中读取，并确保留下构建说明，或者确保它仅使用pyproject.toml构建。

还要注意，GPT-2 的正则表达式在大多数正则表达式引擎中支持不佳，在大多数支持正则表达式的引擎中运行速度也会很慢。我们已验证Oniguruma相当快，并且支持负向预搜索，但如果说 Python 中的正则表达式包有什么不同的话，那就是它甚至更快。

Problem (train_bpe_tinystories): BPE Training on TinyStories#

(a) 在 TinyStories 数据集上训练一个 byte-level BPE tokenizer，最大词汇量为10,000。确保将 TinyStories 的<|endoftext|> special token 添加到词汇表中。将生成的词汇表和合并结果序列化到磁盘上，以供进一步检查。训练耗时多少小时，占用了多少内存？词汇表中最长的标记是什么？它有意义吗？
资源要求：≤30分钟（无GPU），≤30GB 内存 提示：利用预分词过程中的多进程处理以及以下两件事，你应该能够使BPE训练时间低于2分钟：
- (a) <|endoftext|> 标记用于分隔数据文件中的文档。
- (b) 在应用BPE合并之前，<|endoftext|>标记被作为特殊情况处理。
(b) 分析你的代码。在分词器训练过程中，哪一部分耗时最长？

接下来，我们将尝试在 OpenWebText 数据集上训练一个 byte-level BPE tokenizer。与之前一样，我们建议先查看一下该数据集，以便更好地理解其内容。

Problem (train_bpe_expts_owt): BPE Training on OpenWebText#

(a) 在 OpenWebText 数据集上训练一个 byte-level BPE tokenizer，使用最大词汇量为32,000。将生成的词汇表和合并结果序列化到磁盘上以供进一步检查。词汇表中最长的词是什么？它有意义吗？ 资源要求：≤ 12小时（无GPU），≤ 100GB内存
(b) 比较并对比你在TinyStories上训练的分词器与在OpenWebText上训练的 tokenizer。

2.6 BPE Tokenizer: Encoding and Decoding#

在上一部分的任务中，我们实现了一个函数，用于基于输入文本训练一个 BPE 分词器，从而获得分词器**词汇表（vocab）**和 BPE 合并列表（merges）。

现在，我们将实现一个会加载所提供的词汇表（vocab）和合并列表（merges），并使用它们将文本编码为和解码为 token IDs 的BPE分词器。

2.6.1 Encoding text#

通过 BPE 对文本进行编码的过程与我们训练 BPE 词汇表的过程类似，主要有以下几个步骤：

Step 1: Pre-tokenize
我们首先对序列进行预分词处理，并将每个预分词表示为一系列的 UTF-8 字节，就像我们在 BPE training 中所做的那样。我们将把这些字节在每个 pre-token 内部合并为词汇元素，并且会独立处理每个预分词（不在预分词边界处进行合并）。
Step 2: Apply the merges
然后，我们取在 BPE training 过程中生成的词汇元素合并序列，并按照其创建的顺序将其应用于我们的预标记中。

Example(bpe_encoding): BPE encoding example#

假设我们输入字符串 the cat ate，我们的 vocabulary 是：

{0: b' ', 1: b'a', 2: b'c', 3: b'e', 4: b'h', 5: b't', 6: b'th', 7: b' c', 8: b' a', 9: b'the', 10: b' at'}

学习到的 merges 是：

[(b't', b'h'), (b' ', b'c'), (b' ', b'a'), (b'th', b'e'), (b' a', b't')]

首先， pre-tokenizer 将会把输入的字符串分割成 ['the', ' cat', ' ate'] 。然后，我们查看每一个 pre-token，并且应用 BPE merges 进行合并。步骤如下：

第一个 pre-token 'the' 最初表示为[b't', b'h', b'e']。
查看我们的 merges 列表，我们发现第一个适用的 merge 是 (b't', b'h')，并使用该 merge 将 pre-token 转换为 [b'th', b'e']。
然后，我们回到 merges 列表，发现下一个适用的 merge 是 (b'th', b'e')，该 merge 将 pre-token 转换为 [b'the']。
最后，回顾 merges 列表，我们发现该字符串没有更多适用的合并（因为整个 pre-token 已被合并为一个单独的 token），因此我们完成了 BPE merges 的应用。相应的整数序列为[9]。

对剩余的 pre-tokens 重复此过程，我们发现应用 BPE merges 后，前缀标记 ' cat' 被表示为 [b' c', b'a', b't']，变成了整数序列 [7, 1, 5]。应用 BPE merges 后，最终的 pre-token ' ate' 被表示为 [b' at', b'e']，变成了整数序列 [10, 3]。因此，输入字符串的最终编码结果为 [9, 7, 1, 5, 10, 3]。

2.6.1.1 Special tokens#

在编码文本时，您的 tokenizer 应能够正确处理用户定义的special tokens。

2.6.1.2 Memory considerations#

假设我们想要对一个无法装入内存的大型文本文件进行分词。为了高效地对这个大型文件（或任何其他数据流）进行分词，我们需要将其拆分成可管理的块（chunks），并依次处理每个块，这样内存复杂度就会保持恒定，而不是与文本大小呈线性关系。

在此过程中，我们需要确保一个 token 不会跨越块边界，否则我们得到的分词结果将与在内存中对整个序列进行分词的简单方法不同。

2.6.2 Decoding text#

要将一串整数 token IDs 解码回原始文本（raw text），我们只需在词汇表（一个字节序列）中查找每个ID对应的条目，将它们连接在一起，然后将这些字节解码为Unicode字符串。

请注意，输入的 IDs 并不保证能映射到有效的 Unicode 字符串（因为用户可以输入任何整数 IDs 序列）。如果输入的 token IDs 无法生成有效的 Unicode 字符串，则应使用官方的 Unicode 替换字符 U+FFFD 替换格式错误的字节。bytes.decode 的 errors 参数控制着如何处理 Unicode 解码错误，使用 errors='replace' 会自动用替换标记替换格式错误的数据。

Problem (tokenizer): Implementing the tokenizer#

实现一个 Tokenizer 类，该类给定一个 vocabulary 和一组 merges，将文本编码为整数 IDs ，并将整数 IDs 解码为文本。您的 tokenizer 还应支持用户提供的 special tokens（如果它们不在 vocabulary 中，则将其添加到 vocabulary 中）。我们建议使用以下接口：

def __init__(self, vocab, merges, special_tokens=None)
'''
利用给定的 vocabulary、merges 和一组(可选的) special tokens 构造 tokenizer

函数接收以下参数:
    vocab: dict[int, bytes]
    merges: list[tuple[bytes, bytes]]
    special_tokens: list[str] | None = None
'''

def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None)
'''
一个类方法，用于从序列化的 vocabulary 和 merges（格式与BPE training 代码输出相同）以及(可选的) special tokens 中构建并返回一个 Tokenizer

函数接收以下参数:
    vocab_filepath: str
    merges_filepath: str
    special_tokens: list[str] | None = None
'''

def encode(self, text: str) -> list[int]
'''
将输入文本编码为一系列 token IDs
'''

def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]
'''
给定一个字符串的可迭代对象（例如，一个 Python 文件句柄），返回一个生成器，该生成器以惰性方式生成 token IDs。这对于我们无法直接加载到内存中的大型文件的 memory-efficient tokenization 是必需的。
'''

def decode(self, ids: list[int]) -> str
'''
将一系列 token IDs 解码为文本
'''

备注#

要使用我们提供的测试来测试你的 Tokenizer，你首先需要在 [adapters.get_tokenizer] 处实现测试 adapter。然后，运行uv run pytest tests/test_tokenizer.py。你的实现应该能够通过所有测试。

Problem (tokenizer_experiments): Experiments with tokenizers#

(a) 从 TinyStories 和 OpenWebText 中抽取10份文档。使用你之前训练好的 TinyStories 和 OpenWebText 分词器（词汇量分别为10K和32K），将这些抽取的文档编码为整数 IDs。每个分词器的压缩比（字节/词）是多少？
(b) 如果使用 TinyStories 分词器对 OpenWebText 样本进行分词，会发生什么？比较压缩比，定性描述会发生什么。
(c) 估算你的分词器的吞吐量（例如，以字节/秒为单位）。对 Pile 数据集（825GB文本）进行分词需要多长时间？
(d) 使用你的 TinyStories 和 OpenWebText 分词器，将相应的训练数据集和开发数据集编码为一系列整数 token IDs。我们稍后将使用这些 token IDs 来训练我们的语言模型。我们建议将 token IDs 序列化为 uint16 数据类型的 NumPy 数组。为什么选择 uint16 是合适的选择？