Stanford CS336 | Assignment 1 - BPE Tokenizer Training 实现

所有关于 assignment1 的代码已开源在：
https://github.com/ACEEE-1222/Standford-CS336-Assignment-1
如果对你有帮助的话，记得顺手点个star喔！

本文记录了我在 Stanford CS336 第一次作业中完成的一项任务：从头实现一个 字节级 BPE（Byte Pair Encoding）分词器训练器，并支持高效的并行预分词与合并操作优化。本文将介绍 BPE 分词器的实现细节、并行预处理策略、特殊符号的处理方式，以及对 merge 步骤的优化。

一、作业要求概述

本次任务目标是在 TinyStories 数据集上训练一个字节级 BPE 分词器。如图所示，需要完成函数train_bpe。

核心要求如下：

1. 初始化：以所有可能的字节（0-255）作为初始词汇表；
2. 预处理：将文本分割为初始令牌（字节序列）；
3. 合并：统计所有字节对的频率，合并最频繁的对并更新词汇表，重复此过程直至达到目标词汇量。

二、整体实现结构

整个分词器训练流程主要分为以下三步：

1. 初始化词表

• 首先将所有 ASCII 字节值（0~255）加入初始词表。
• 然后添加用户指定的 special_tokens，如 <|endoftext|>、<pad> 等。

# 1. Vocabulary Initialization
vocab = {i: bytes([i]) for i in range(256)}
for tok in special_tokens:
    vocab[len(vocab)] = tok.encode("utf-8")

2. 并行预分词（Pre-tokenization）

（1）文档边界切分

为了并行处理，我们需要将文件划分为若干个 chunk，每个 chunk 的起始位置应当正好落在一个 <|endoftext|> 标记上，确保不会跨文档切分。我们使用 find_chunk_boundaries 函数来实现这一功能：

boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")

该函数会从初始估计位置开始，向前读取直到找到一个 <|endoftext|> 再作为真正边界。

（2）分发任务到子进程

我们利用 Python 的 multiprocessing.Pool，为每个 chunk 分发 process_chunk 任务：

task_args = [(input_path, start, end, special_tokens) for start, end in zip(boundaries[:-1], boundaries[1:])]
with Pool(processes=num_processes) as pool:
    chunk_results = pool.map(process_chunk, task_args)

每个 chunk 的处理逻辑包括：

• 用正则表达式 re.split 去除所有特殊 token，并按文档分别处理；
• 使用 GPT-2 的正则表达式模式进行预分词；
• 每个 token 转为字节流，再按字节分割成 [b1, b2, ...] 的形式。

3. BPE 合并训练

训练阶段按照以下流程进行：

• 初始化 counts 和 pair_to_indices，分别用于记录 pair 频率和其在哪些 token 中出现。
• 不断选择出现频率最高的 pair (a, b) 进行合并，生成新 token。
• 将受影响的所有 token 中的旧 pair 信息删除，并插入新 pair。
• 直到词表大小达到目标 vocab_size。

优化关键点在于：

• 只更新受到影响的 pair 计数，而不是每次全局重新统计；
• 避免在迭代过程中修改原集合，采用 .copy() 确保稳定性。

示例代码片段（合并过程）：

需要注意的是：在 affected_indices 用 pair_to_indices[max_pair].copy() 创建避免集合被动态修改。以及我一开始将 for j in affected_indices 写成三次遍历，导致pair_to_indices[max_pair]中间状态交叉污染，一次遍历就不会存在这样的问题（可以想想为什么）。

# 3. Compute BPE merges
    merges : list[tuple[bytes, bytes]] = []
    pre_tokens_bytes: list[list[bytes]] = [token for chunk in chunk_results for token in chunk]
    counts = defaultdict(int)
    pair_to_indices = defaultdict(set)
    for idx, token in enumerate(pre_tokens_bytes):
        for i in range(len(token) - 1):
            pair = (token[i], token[i + 1])
            counts[pair] += 1
            pair_to_indices[pair].add(idx)

    idx = len(vocab)
    while idx < vocab_size:
        if not counts:
            break
            
        max_pair: tuple[bytes, bytes] = None
        max_cnt= -1
        for pair, cnt in counts.items():
            if cnt > max_cnt:
                max_pair = pair
                max_cnt = cnt
            elif cnt == max_cnt:
                if max_pair is None or pair > max_pair:
                    max_pair = pair

        merges.append(max_pair)
        a, b = max_pair
        new_token = a + b
        vocab[idx] = new_token
        idx += 1

        affected_indices = pair_to_indices[max_pair].copy()
        for j in affected_indices:
            token = pre_tokens_bytes[j]
            for i in range(len(token) - 1):
                old_pair = (token[i], token[i+1])
                pair_to_indices[old_pair].discard(j)
                counts[old_pair] -= 1
                if counts[old_pair] == 0:
                    counts.pop(old_pair)
                    pair_to_indices.pop(old_pair, None)

            merged = []
            i = 0
            while i < len(token):
                if i < len(token) - 1 and token[i] == a and token[i+1]==b:
                    merged.append(new_token)
                    i += 2
                else:
                    merged.append(token[i])
                    i += 1
            pre_tokens_bytes[j]=merged

            token = pre_tokens_bytes[j]
            for i in range(len(token) - 1):
                pair = (token[i], token[i + 1])
                counts[pair] += 1
                pair_to_indices[pair].add(j)

---

三、关键模块代码解释

1. find_chunk_boundaries

用于在特殊 token 边界上切割数据块，避免语义跨越：

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.
    """
    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:
                true_position = initial_position + found_at
                chunk_boundaries[bi] = true_position
                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))

2. process_chunk

完成：

• 特殊 token 切割；
• 分词（正则）；
• 字节分解为 [b1, b2, ...]：

def process_chunk(args: tuple[str, int, int, list[str]]) -> list[list[bytes]]:
    input_path, start, end, special_tokens = args
    """
    Processes a chunk of the input file and returns byte pair frequency counts.

    Args:
        input_path (str): The path of input file.
        start (int): Start byte offset of the chunk.
        end (int): End byte offset of the chunk.
        special_tokens (list[str]): List of special tokens that should not be merged across.

    Returns:
        pre_token_bytes (list[list[bytes]]): list of tokens, where each token is a list of bytes
    """

    with open(input_path, "rb") as file:
        file.seek(start)
        chunk = file.read(end - start).decode("utf-8", errors="ignore")

    # 1. Remove special tokens by splitting the chunk at those tokens
    pattern = "|".join(re.escape(tok) for tok in special_tokens)
    documents = re.split(pattern, chunk)

    # 2. Pre-tokenize and count byte pair frequencies
    pre_tokens_bytes: list[list[bytes]] = []
    PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
    for doc in documents:
        tokens = [match.group(0).encode("utf-8") for match in re.finditer(PAT, doc)]
        for token in tokens:
            token_bytes = [bytes([b]) for b in token]
            pre_tokens_bytes.append(token_bytes)

    return pre_tokens_bytes

---

四、最终结果

最终返回两个结构：

vocab: dict[int, bytes]  # token id -> bytes
merges: list[tuple[bytes, bytes]]  # 按顺序记录合并对

在 tests/adapters.py 的 run_train_bpe 中调用函数：

    from cs336_basics.bpe_tokenizer import train_bpe

    return train_bpe(input_path,vocab_size,special_tokens)

---

五、测试通过情况与调试经验

作业附带三个测试用例，涵盖了基本词表构造、特殊 token 行为，以及合并顺序的正确性。
用 uv run pytest tests/test_train_bpe.py 命令进行测试，结果如下：

六、完整代码（bpe_tokenizer.py）

import os
import regex as re
from typing import BinaryIO
from multiprocessing import Pool
from collections import defaultdict

def train_bpe(
    input_path: str,
    vocab_size: int,
    special_tokens: list[str],
    num_processes: int = 8
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """
    Trains a byte-level BPE (Byte Pair Encoding) tokenizer on the given input text file.

    Parameters
    ----------
    input_path : str
        Path to a UTF-8 encoded text file containing training data for the BPE tokenizer.
        Each line is considered part of the corpus.

    vocab_size : int
        The total size of the final vocabulary (must include initial byte-level tokens,
        all merged tokens produced during training, and the given special tokens).

    special_tokens : list[str]
        A list of user-defined special tokens (e.g., ["<|endoftext|>", "<pad>"]) to be 
        added to the vocabulary. These tokens do NOT participate in merge decisions.

    num_processes : int, optional (default=8)
        Number of parallel processes used during pre-tokenization. Each process handles
        a chunk of the input corpus split at special token boundaries. More processes
        generally mean faster pre-tokenization.

    Returns
    -------
    vocab : dict[int, bytes]
        A dictionary mapping token IDs (integers) to token values (in bytes). The token 
        IDs should be assigned sequentially starting from 0.

    merges : list[tuple[bytes, bytes]]
        A list of BPE merge operations, where each tuple represents two byte-level tokens 
        that were merged together. The list should be ordered by merge time (first merge first).
    """

    # 1. Vocabulary Initialization
    vocab = {i: bytes([i]) for i in range(256)}
    for tok in special_tokens:
        vocab[len(vocab)] = tok.encode("utf-8")


    # 2. Pre-tokenization
    with open(input_path, "rb") as f:
        boundaries = find_chunk_boundaries(f, num_processes, "<|endoftext|>".encode("utf-8"))

    task_args = [(input_path, start, end, special_tokens) for start, end in zip(boundaries[:-1], boundaries[1:])]
    with Pool(processes=num_processes) as pool:
        chunk_results = pool.map(process_chunk, task_args)
    
    # 3. Compute BPE merges
    merges : list[tuple[bytes, bytes]] = []
    pre_tokens_bytes: list[list[bytes]] = [token for chunk in chunk_results for token in chunk]
    counts = defaultdict(int)
    pair_to_indices = defaultdict(set)
    for idx, token in enumerate(pre_tokens_bytes):
        for i in range(len(token) - 1):
            pair = (token[i], token[i + 1])
            counts[pair] += 1
            pair_to_indices[pair].add(idx)

    idx = len(vocab)
    while idx < vocab_size:
        if not counts:
            break
            
        max_pair: tuple[bytes, bytes] = None
        max_cnt= -1
        for pair, cnt in counts.items():
            if cnt > max_cnt:
                max_pair = pair
                max_cnt = cnt
            elif cnt == max_cnt:
                if max_pair is None or pair > max_pair:
                    max_pair = pair

        merges.append(max_pair)
        a, b = max_pair
        new_token = a + b
        vocab[idx] = new_token
        idx += 1

        affected_indices = pair_to_indices[max_pair].copy()
        for j in affected_indices:
            token = pre_tokens_bytes[j]
            for i in range(len(token) - 1):
                old_pair = (token[i], token[i+1])
                pair_to_indices[old_pair].discard(j)
                counts[old_pair] -= 1
                if counts[old_pair] == 0:
                    counts.pop(old_pair)
                    pair_to_indices.pop(old_pair, None)

            merged = []
            i = 0
            while i < len(token):
                if i < len(token) - 1 and token[i] == a and token[i+1]==b:
                    merged.append(new_token)
                    i += 2
                else:
                    merged.append(token[i])
                    i += 1
            pre_tokens_bytes[j]=merged

            token = pre_tokens_bytes[j]
            for i in range(len(token) - 1):
                pair = (token[i], token[i + 1])
                counts[pair] += 1
                pair_to_indices[pair].add(j)

    return vocab, merges

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.
    """
    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:
                true_position = initial_position + found_at
                chunk_boundaries[bi] = true_position
                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 process_chunk(args: tuple[str, int, int, list[str]]) -> list[list[bytes]]:
    input_path, start, end, special_tokens = args
    """
    Processes a chunk of the input file and returns byte pair frequency counts.

    Args:
        input_path (str): The path of input file.
        start (int): Start byte offset of the chunk.
        end (int): End byte offset of the chunk.
        special_tokens (list[str]): List of special tokens that should not be merged across.

    Returns:
        pre_token_bytes (list[list[bytes]]): list of tokens, where each token is a list of bytes
    """

    with open(input_path, "rb") as file:
        file.seek(start)
        chunk = file.read(end - start).decode("utf-8", errors="ignore")

    # 1. Remove special tokens by splitting the chunk at those tokens
    pattern = "|".join(re.escape(tok) for tok in special_tokens)
    documents = re.split(pattern, chunk)

    # 2. Pre-tokenize and count byte pair frequencies
    pre_tokens_bytes: list[list[bytes]] = []
    PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
    for doc in documents:
        tokens = [match.group(0).encode("utf-8") for match in re.finditer(PAT, doc)]
        for token in tokens:
            token_bytes = [bytes([b]) for b in token]
            pre_tokens_bytes.append(token_bytes)

    return pre_tokens_bytes

---

如有问题欢迎留言讨论。