本文是 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+""")

简而言之这个正则会将语料库按照空格拆分，并且考虑到了很多边缘情况。

我们来看一下这个正则的切分结果：

>>> 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")