本文是 Stanford CS336 课程学习记录的一部分。
概述
在语言建模中，分词（Tokenization）的作用是将连续的自然语言分割成机器可处理的一连串离散的小单元（token），并为每个 token 分配唯一的 ID，便于后续的计算。如果以整句为单位，那么自然语言中的句子数量几乎无穷无尽，模型不可能认识所有句子。通过拆分成 Token，模型可以利用有限的词表组合出无穷的表达。好的分词算法能将从未见过的词（如新造的合成词）拆解为已知的部分，避免模型出现“看不懂”的情况。
BPE
BPE 是一种比较主流的子词级别的分词算法，它的基本思想是将语料拆分成字节，然后逐步合并出现频次最多的字节对，直到词表大小达到规定值。
BPE 分词器的训练可以大致分为三个步骤：
- 初始化词表。由于我们训练的是一个字节级别的 BPE 分词器，初始的词表就是所有可能的 256 种字节（一个字节有 8 位，2^8=256）。
- 预分词。预分词的主要作用是切分词的边界，避免后续合并时跨词边界合并，产生无意义的 token。
据说还有加速的作用，但是我没看懂，代码也没管预分词加速这一块。
- 执行合并并记录合并过程。在预分词得到的每个词的内部进行统计，找出整个语料库上出现次数最多的相邻 token 对，然后将这个 token 对合并为一个新的 token 并更新计数。将新的 token 加入词表，并将执行的合并操作记录到一个合并日志中。重复上述过程，直到词表大小达到预定的值。
预分词
对于字母语言，我们通常在空格处对词进行切分就可以了。但实际自然语言中有很多复杂的模式，需要谨慎处理。GPT-2 使用了一个复杂的正则表达式来完成预分词过程：
PAT: re.Pattern[str] = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")Qwen3.5-Plus 对这个正则的分析
该正则表达式主要用于自然语言处理（NLP）任务中的文本分词（Tokenization）。具体而言，它是 GPT-2 及其衍生模型（如某些版本的 BPE 分词器）中常见的预处理模式，旨在将文本分割为具有语义意义的单元，同时保留空格和标点符号的结构信息。以下是针对该正则表达式的详细技术分析：
1. 技术前提与兼容性
值得注意的是，标准 Python 库 re 并不支持 \p{L} 和 \p{N} 这样的 Unicode 属性语法。
- 依赖库：该表达式通常需要配合第三方库
regex（PyPI 包名）使用，而非标准的re模块。regex模块提供了对 Unicode 属性的完整支持。 - 代码上下文：虽然代码片段中使用了
re.compile，但在实际运行环境中，re可能被别名指向regex模块，或者该代码片段依赖于支持 Unicode 属性的特定环境配置。
2. 表达式结构分解
该正则表达式由多个通过竖线 | 分隔的分支组成。正则引擎会按照从左到右的顺序尝试匹配，一旦匹配成功，即停止后续分支的尝试。
-
'(?:[sdmt]|ll|ve|re)- 功能：匹配常见的英语缩写形式。
- 说明：匹配单引号
'后接特定后缀（s,d,m,t,ll,ve,re）。 - 示例：
's(it’s),'ll(we’ll),'ve(I’ve)。 - 目的：防止缩写被错误地分割为单词和标点符号，保持语义完整性。
-
?\p{L}+- 功能：匹配单词（字母序列）。
- 说明：
?表示可选的前导空格（0 或 1 个），\p{L}+匹配一个或多个 Unicode 字母字符。 - 目的：识别文本中的词汇单元，支持多语言字符。
-
?\p{N}+- 功能：匹配数字序列。
- 说明：
?表示可选的前导空格，\p{N}+匹配一个或多个 Unicode 数字字符。 - 目的：将数值识别为独立的 token。
-
?[^\s\p{L}\p{N}]+- 功能：匹配标点符号或特殊符号。
- 说明：
[^\s\p{L}\p{N}]+匹配任何非空白、非字母、非数字的字符序列。 - 目的：捕获 punctuation 和其他特殊字符。
-
\s+(?!\S)- 功能：匹配末尾或独立的空白字符。
- 说明：
\s+匹配一个或多个空白字符，(?!\S)是负向先行断言，确保空白字符后不跟随非空白字符（即后面是字符串结尾或仍是空白）。 - 目的：专门处理文本末尾的空白或连续的空白块，确保空格信息不被丢失，这对于基于字节对编码（BPE）的分词器重建文本至关重要。
-
\s+- 功能：匹配剩余的空白字符。
- 说明：作为捕获空白字符的兜底规则。
- 目的：确保所有的空白字符都能被 token 化流程捕获。
3. 逻辑流程总结
该正则表达式的匹配优先级设计体现了以下逻辑：
- 优先处理特殊结构：首先检查是否为英语缩写，避免拆分。
- 内容优先：其次匹配具体的语义内容（单词、数字、符号），并尝试吸收其前导空格。
- 空白处理：最后处理空白字符，区分普通空白和末尾空白，以确保分词后的文本能够精确还原原始格式。
4. 应用场景
- 大型语言模型（LLM）预处理：此模式广泛用于 GPT-2、GPT-3 等模型的 Tokenizer 实现中（例如 Hugging Face
transformers库中的GPT2Tokenizer）。 - 多语言支持：由于使用了 Unicode 属性（
\p{L},\p{N}），该表达式不仅适用于英语，也能正确处理中文、日文等其他语言的字符分类。 - 数据清洗与分割：用于需要将文本流精确分割为语义单元且需保留空格信息的 NLP 管道中。
结论
该正则表达式是一个高度专业化的分词器模式（Tokenizer Pattern）。其核心作用是将原始文本流分割为单词、数字、标点、缩写及空白符序列，同时利用 Unicode 属性确保多语言兼容性，并利用特定的空白匹配逻辑保障文本的可逆性（Reversibility）。在使用时，请确保环境已安装并正确引入了 regex 库。
简而言之这个正则会将语料库按照空格拆分，并且考虑到了很多边缘情况。
我们来看一下这个正则的切分结果：
>>> import regex as re
>>> re.findall(PAT, "You are a helpful assistant.")
['You', ' are', ' a', ' helpful', ' assistant', '.']合并
预分词之后，我们得到了一系列预分词的 token，其中每个 token 都可用一个 UTF-8 字节序列来表示。在合并过程中，我们统计语料库中出现次数最多的相邻字节对，然后将这两个字节合并成一个新的 token，然后将这个新的 token 添加到词表。这样训练得到的词表的大小就等于初始的 256 个字节，加上我们执行的合并的总次数。
如果统计时得到了若干个 token 对同时具有最多的出现次数，那么我们取每个 token 对的第一个 token 进行字典序（lexicographic order）排序，取最大者作为合并的对象。
字典序通常用于比较字符串。其规则是：从第一个字符开始比较，每个字符都对应了 Unicode 中的一个码点，码点较大者字典序较大，例如 B 大于 A。如果当前字符相同，则比较下一个字符。如果前缀的所有字符都相同，则总长度较长的那个字符串字典序较大。
max([("A", "B"), ("A", "C"), ("B", "ZZ"), ("BA", "A")])
('BA', 'A')我们需要注意一下特殊 token。有一类 token 是我们规定用于特殊用途的，例如，标记文本的结束。这一类 token 不应该参与合并过程，并且应该在合并之前添加到初始化的词表。在作业中，我们使用 <|endoftext|> 作为特殊 token。
实现
按照作业指导的要求，我们应该实现一个 BPE 分词器的训练代码，它接受以下输入：
input_path：语料库文件的路径vocab_size：预设的词表大小special_tokens：特殊 token 列表
然后返回以下内容：
vocab：词表，用一个dict[int, bytes]来表示merges：合并记录，用一个list[tuple[bytes, bytes]]来表示。这个列表的顺序就是产生合并的顺序。
完整代码如下
import os
from multiprocessing import Pool
from typing import BinaryIO
import regex as re
from tqdm import tqdm, trange
# 编译正则表达式以提高效率
# GPT-2 预分词正则表达式
PAT: re.Pattern[str] = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
def find_chunk_boundaries(file: BinaryIO, desired_num_chunks: int, split_special_token: bytes) -> list[int]:
"""
Chunk the file into parts that can be counted independently.
May return fewer chunks if the boundaries end up overlapping.
This function is copied from pretokenization_example.py and the file is deleted now.
"""
assert isinstance(split_special_token, bytes), "Must represent special token as a bytestring"
# Get total file size in bytes
file.seek(0, os.SEEK_END)
file_size = file.tell()
file.seek(0)
chunk_size = file_size // desired_num_chunks
# Initial guesses for chunk boundary locations, uniformly spaced
# Chunks start on previous index, don't include last index
chunk_boundaries = [i * chunk_size for i in range(desired_num_chunks + 1)]
chunk_boundaries[-1] = file_size
mini_chunk_size = 4096 # Read ahead by 4k bytes at a time
for bi in range(1, len(chunk_boundaries) - 1):
initial_position = chunk_boundaries[bi]
file.seek(initial_position) # Start at boundary guess
while True:
mini_chunk = file.read(mini_chunk_size) # Read a mini chunk
# If EOF, this boundary should be at the end of the file
if mini_chunk == b"":
chunk_boundaries[bi] = file_size
break
# Find the special token in the mini chunk
found_at = mini_chunk.find(split_special_token)
if found_at != -1:
chunk_boundaries[bi] = initial_position + found_at
break
initial_position += mini_chunk_size
# Make sure all boundaries are unique, but might be fewer than desired_num_chunks
return sorted(set(chunk_boundaries))
def split_special_tokens(text: str, special_tokens: list[str]) -> list[str]:
"""
根据 special tokens 分割文本，并返回纯文本 chunk 列表。
这些 chunk 可以安全地进行后续的预分词。
"""
if not special_tokens:
return [text]
# 1. 转义所有特殊标记
# re.escape() 会处理特殊字符，例如 <|endoftext|> 会变成 <\|endoftext\|>
escaped_tokens = [re.escape(token) for token in special_tokens]
# 2. 构建正则表达式模式：使用 '|' 连接所有转义后的标记
# 例如：<\|endoftext\|>|\[CLS\]|\[SEP\]
delimiter_pattern = "|".join(escaped_tokens)
# 3. 使用 re.split 分割文本
# re.split 会移除匹配到的分隔符，并返回分隔符之间的内容
raw_chunks = re.split(delimiter_pattern, text)
# 4. 过滤掉可能出现的空字符串（例如，如果分隔符在开头、结尾或连续出现）
# 示例: "TOKEN1<SP_TOKEN>TOKEN2" -> ['TOKEN1', 'TOKEN2']
# 示例: "<SP_TOKEN>TOKEN1<SP_TOKEN>TOKEN2" -> ['', 'TOKEN1', 'TOKEN2'] -> ['TOKEN1', 'TOKEN2']
# 示例: "TOKEN1<SP_TOKEN>" -> ['TOKEN1', ''] -> ['TOKEN1']
text_chunks = [chunk for chunk in raw_chunks if chunk]
return text_chunks
def pre_tokenize_for_chunk(args: tuple[str | os.PathLike, list[str], int, int]) -> dict[tuple[bytes, ...], int]:
pre_token_counts: dict[tuple[bytes, ...], int] = {}
input_path, special_tokens, start, end = args
with open(input_path, mode="rb") as f:
f.seek(start)
# 截取当前线程的 chunk
chunk = f.read(end - start).decode("utf-8", errors="ignore")
# 根据 special tokens 进行分割，确保 special tokens 被移除且这些 token 两侧的文本不会被粘连
splitted_chunks = split_special_tokens(chunk, special_tokens)
# 针对分割出的每个 chunk 进行预分词
for splitted_chunk in splitted_chunks:
for match in PAT.finditer(splitted_chunk):
pre_token = match.group(0)
byte_tokens = tuple(bytes([b]) for b in pre_token.encode("utf-8"))
pre_token_counts[byte_tokens] = pre_token_counts.get(byte_tokens, 0) + 1
return pre_token_counts
def pre_tokenize(
input_path: str | os.PathLike,
special_tokens: list[str],
) -> dict[tuple[bytes, ...], int]:
pre_token_counts: dict[tuple[bytes, ...], int] = {}
# 一次性读取文件并确定边界
with open(input_path, mode="rb") as f:
# 根据CPU核心数确定分块数量，但设置上限避免过多小任务的开销
cpu_count = os.cpu_count()
if not cpu_count:
raise ValueError("indeterminable CPU count")
num_chunks = min(cpu_count * 2, 16)
boundaries = find_chunk_boundaries(f, num_chunks, "<|endoftext|>".encode("utf-8")) # noqa: UP012
args = []
for start, end in zip(boundaries[:-1], boundaries[1:]):
args.append((input_path, special_tokens, start, end))
with Pool(cpu_count) as pool:
for chunk_result in tqdm(
pool.imap_unordered(pre_tokenize_for_chunk, args), total=len(args), desc="pre-tokenize chunks", unit="chunk"
):
for token, count in chunk_result.items():
pre_token_counts[token] = pre_token_counts.get(token, 0) + count
return pre_token_counts
def get_pair_counts(words: dict[tuple[bytes, ...], int]) -> dict[tuple[bytes, bytes], int]:
"""
统计所有相邻 token 对的频率
"""
pair_counter: dict[tuple[bytes, bytes], int] = {}
for word, frequency in words.items():
word_len = len(word)
# 遍历每个词中相邻的 token 对
for i in range(word_len - 1):
pair = (word[i], word[i + 1])
# 累加该对的出现次数
pair_counter[pair] = pair_counter.get(pair, 0) + frequency # {(b'h', b'e'): 5, (b'l', b'l'): 10, ...}
return pair_counter
def merge_pair(
words: dict[tuple[bytes, ...], int], target_pair: tuple[bytes, bytes], pair_counts: dict[tuple[bytes, bytes], int]
) -> dict[tuple[bytes, ...], int]:
new_words: dict[tuple[bytes, ...], int] = {}
new_token = target_pair[0] + target_pair[1]
for word, freq in words.items():
if len(word) < 2:
new_words[word] = freq
continue
# Check if the target_pair is in the word
if target_pair not in zip(word, word[1:]):
new_words[word] = freq
continue
# 1. Deduct old pair counts
for i in range(len(word) - 1):
pair = (word[i], word[i + 1])
pair_counts[pair] -= freq
if pair_counts[pair] <= 0:
pair_counts.pop(pair, None)
# 2. Form new word
new_word_list = []
i = 0
while i < len(word):
if i < len(word) - 1 and (word[i], word[i + 1]) == target_pair:
new_word_list.append(new_token)
i += 2
else:
new_word_list.append(word[i])
i += 1
new_word = tuple(new_word_list)
new_words[new_word] = new_words.get(new_word, 0) + freq
# 3. Add new pair counts
for i in range(len(new_word) - 1):
pair = (new_word[i], new_word[i + 1])
pair_counts[pair] = pair_counts.get(pair, 0) + freq
return new_words
def get_best_pair(pair_counts: dict[tuple[bytes, bytes], int]) -> tuple[bytes, bytes]:
"""
获取最高频的 token pair
- 如果最高频的只有一个，则返回
- 如果有多个，则按照字典顺序排序
"""
return max(pair_counts.keys(), key=lambda p: (pair_counts[p], p[0], p[1]))
def train_bpe(
input_path: str | os.PathLike,
vocab_size: int,
special_tokens: list[str],
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
# 生成初始包含 256 个字节字符
vocab: list[bytes] = [bytes([i]) for i in range(0, 256)]
# 添加 special_tokens
vocab += [special_token.encode("utf-8") for special_token in special_tokens]
# 预分词
words = pre_tokenize(input_path, special_tokens)
merges: list[tuple[bytes, bytes]] = []
# 获取每个词中相邻的 token 对出现的次数
# 仅在初始化时全量统计一次
pair_counts = get_pair_counts(words)
initial_vocab_len = len(vocab)
# 计划新增 token 数
planned_new = max(0, vocab_size - initial_vocab_len)
# 使用 trange 显示合并进度；如果中途 pair_counts 为空会提前中断
for _ in trange(planned_new, desc="BPE merges", unit="merge"):
# 空 dict 将被视为 False，这表示语料库中所有的词都无法再分割出 token 了
if not pair_counts:
break
# 获取需要合并的下一个 pair
best_pair = get_best_pair(pair_counts)
# 合并产生新 token
new_token = best_pair[0] + best_pair[1]
# 将新 token 加入词表和集合
vocab.append(new_token)
# 将本次的 pair 加入合并记录
merges.append(best_pair)
words = merge_pair(words, best_pair, pair_counts)
print(f"Final vocab size: {len(vocab)}")
return {i: v for i, v in enumerate(vocab)}, merges这是一个欠优化的代码。虽然能够通过作业中的速度测试，但是仍然缺少了很多关键的速度优化。
编码和解码
from __future__ import annotations
import json
from collections.abc import Iterable, Iterator
from typing import cast
import regex as re
from cs336_basics.bpe import PAT
def pre_tokenize(text: str, special_tokens: list[str] | None = None) -> list[list[bytes]]:
"""
执行预分词过程。注意这里和 bpe.py 中不同，不能排除 special_tokens
"""
if special_tokens:
# 为了让正则优先匹配更长的 special token（避免被其子串先匹配），按长度降序排序
sorted_specials = cast(list[str], sorted(special_tokens, key=len, reverse=True))
escaped = [re.escape(s) for s in sorted_specials]
pattern = "(" + "|".join(escaped) + ")"
parts = re.split(pattern, text)
special_set = set(special_tokens)
else:
parts = [text]
special_set = set()
pre_tokens: list[list[bytes]] = []
for part in parts:
if not part:
continue
# 如果 part 是 exact match 的 special token，则把整个 token 当成一个 bytes 单元（不拆字节）
if part in special_set:
pre_tokens.append([part.encode("utf-8")])
else:
for match in PAT.finditer(part):
pre_token = match.group(0)
byte_tokens = [bytes([b]) for b in pre_token.encode("utf-8")]
pre_tokens.append(byte_tokens)
return pre_tokens
class Tokenizer:
def __init__(
self, vocab: dict[int, bytes], merges: list[tuple[bytes, bytes]], special_tokens: list[str] | None = None
):
"""
Args:
vocab: dict[int, bytes]
merges: list[tuple[bytes, bytes]]
special_tokens: list[str] | None = None
"""
self.vocab = vocab
self.merges = merges
self.special_tokens = special_tokens
# Build reverse vocab mapping bytes -> id for fast lookup
self._reverse_vocab: dict[bytes, int] = {v: k for k, v in self.vocab.items()}
self._merge_ranks: dict[tuple[bytes, bytes], int] = {}
for idx, pair in enumerate(self.merges):
self._merge_ranks[pair] = idx
@staticmethod
def from_files(vocab_filepath: str, merges_filepath: str, special_tokens: list[str] | None = None) -> Tokenizer:
"""
Args:
vocab_filepath: str
merges_filepath: str
special_tokens: list[str] | None = None
"""
with open(vocab_filepath, encoding="utf-8") as f:
vocab_data = json.load(f)
vocab = {v: k.encode("latin-1") for k, v in vocab_data.items()}
merges = []
with open(merges_filepath, encoding="utf-8") as f:
for merge in f:
parts = merge.strip().split()
if len(parts) == 2:
merges.append((parts[0].encode("latin-1"), parts[1].encode("latin-1")))
return Tokenizer(vocab, merges, special_tokens)
def encode(self, text: str) -> list[int]:
# Step 1: pre-tokenize
pre_tokens = pre_tokenize(text, self.special_tokens)
ids: list[int] = []
for curr_token in pre_tokens:
# curr_token: list[bytes], 每个 bytes 是单个字节字符（b'a' 等）
# 使用 merges rank 来进行贪心合并：每次只在当前 token 的相邻 pairs 中
# 选择优先级最高（rank 最小）的 pair 合并直到没有可合并对为止。
if len(curr_token) == 0:
continue
# 转为可变列表
token_list = list(curr_token)
while True:
# 生成所有相邻对并找出在 merge_rules 中优先级最小的那个
best_idx = -1
best_rank = None
# pairs_count = len(token_list) - 1
for i in range(len(token_list) - 1):
pair = (token_list[i], token_list[i + 1])
rank = self._merge_ranks.get(pair)
if rank is not None:
if best_rank is None or rank < best_rank:
best_rank = rank
best_idx = i
if best_idx == -1:
break
# 合并 best_idx 和 best_idx+1
new_token = token_list[best_idx] + token_list[best_idx + 1]
token_list[best_idx : best_idx + 2] = [new_token]
# 将 token_list 转换为 ids（使用已缓存的 reverse vocab）
for token in token_list:
if token in self._reverse_vocab:
ids.append(self._reverse_vocab[token])
else:
# fallback：把 token 拆回单字节
# 理论上这不应该发生，只要 vocab 和 merges 是一致的
for b in token:
single = bytes([b])
if single in self._reverse_vocab:
ids.append(self._reverse_vocab[single])
else:
raise KeyError(f"Byte token {single!r} not found in vocab")
return ids
def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
for text in iterable:
yield from self.encode(text)
def decode(self, ids: list[int]) -> str:
decode_bytes = b""
for id in ids:
decode_bytes += self.vocab[id]
# auto replace the malformed bytes with the official Unicode replacement character U+FFFD
return decode_bytes.decode("utf-8", errors="replace")