Stanford CS336 Assignment 1 Section 2: BPEトークナイザー

インドア系ならトランスフォーマー

インドア派の間ではみんな生成AIで何か作っているらしい。最近も淫夢関連で話題の作品が出ていた。

日本発のコンテンツはやはり淫夢が強い。以前の YAJU&U もそうだ。¹ 技術の変化は淫夢を見ればわかる。

ユーザーとして楽しむのもいいが、中身の仕組みが気になるのがインドア系だ。そんなとき、Xである講義を見つけた。

Stanford CS336

https://cs336.stanford.edu/

スタンフォード大学が公開している、LLMをゼロから作る講義 (CS336: Language Modeling from Scratch) だ。PyTorchは使えるが、便利機能は使ってはいけないことになっているらしい。課題コードも https://github.com/stanford-cs336 で公開されている。

各リポジトリの AGENTS.md には、自力での理解を深めさせるためAIによるコード生成や課題解決を禁止する方針が明記されている。しかし、

履修者ではなくTAも活用できないので自己判断でAIを活用するのも1つの手だと思った
AIを使わず時間を掛けすぎて結局最後までできない方がよっぽど学習効果が低く本末転倒である

ことからAIも活用しつつ理解を深め、課題を進めていこうと思う。

全5課題のうち、今回は課題1のセクション2 (Section 2) を取り上げる。このセクションではトークナイザーを実装する。

連載記事になる予定だ。²

トークナイザーとは何か

始める前に、今回のセクションで扱うトークナイザーについて速習する。

トークナイザーは、テキストとLLMが扱う数値 (ID) 列の間を相互に変換する仕組みだ。文章をトークンと呼ばれる単位に分割し、それぞれを語彙表 (vocabulary) に登録されたIDへ置き換える。たとえば 猫が好き は、トークナイザーを通すことで [123, 456, 789] のような整数列になり、LLMはこのID列を入力として受け取って処理を行う。

トークナイザーによる変換 説明用の語彙を使った例

文章

猫が好き

トークンID列

123 456 789

中身は結合済みトークンごとに語彙表を引いている。

猫 E7 8C AB 123

が E3 81 8C 456

好き E5 A5 BD E3 81 8D 789

語彙表として見る

トークン	UTF-8バイト列	語彙ID
猫	`E7 8C AB`	`123`
が	`E3 81 8C`	`456`
好き	`E5 A5 BD E3 81 8D`	`789`

分割して「1文字=1トークン」にすれば未知語はほぼ出ないが入力トークン列が長くなり、「1単語=1トークン」にすれば短くなるが語彙表にない未知語 (OOV: Out-of-Vocabulary) を扱えない。 バイトレベルBPE (Byte Pair Encoding) はテキストを UTF-8 のバイト列として扱い、頻繁に隣り合うバイトペアを繰り返し結合していく手法だ。全バイト (0 - 255) を初期語彙に含むため未知語を原理的に防ぎつつ、頻出表現をまとめてトークン列を短縮できる。³

たとえば low low lower というコーパスなら、頻出する l と o を結合して lo を作り、さらに lo と w から low を作る。BPEの訓練では、この「どのペアをどの順で結合したか」の履歴である結合規則 (merges) と、結合後のバイト列とIDの対応表である語彙表 (vocab) を構築する。エンコード時にも全く同じ順序で規則を適用していく。

`train_bpe`の全体像

train_bpe はテキストのエンコードではなく、コーパスを読み込んで語彙表と結合規則を構築する関数だ。

work/assignment1-basics/cs336_basics/my_tokenizer.pyd8a9247

def train_bpe(
    input_path: str | os.PathLike[str],
    vocab_size: int,
    special_tokens: list[str],
    max_bytes: int | None = None,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """入力コーパスから BPE トークナイザーを訓練する。

    流れは次の 3 段である。
    1. 初期語彙を作る (全バイト + 特殊トークン)。
    2. テキストを pre-token に分けて、各部品の出現回数を数える。
    3. よく隣り合う部品同士を結合し、語彙を目標サイズまで育てる。

    Args:
        input_path (str | os.PathLike[str]): 訓練データのパス。
        vocab_size (int): 語彙サイズ (特殊トークンを含む)。
        special_tokens (list[str]): 語彙に追加する特殊トークン。これらの文字列は
            複数トークンに分割されず、常に単一トークンとして扱う。
        max_bytes (int | None): 先頭から使う最大バイト数。None ならファイル全体。
            公式テストは渡さないので、提出用の挙動は変わらない。

    Returns:
        tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
            vocab: 訓練後の語彙。トークン ID からバイト列への写像。
            merges: 作成順の BPE マージ。各要素は (token1, token2) で、
                token1 が token2 とマージされたことを表す。
    """
    if debug_on():
        print("[1/3] initialize vocabulary", flush=True)
    vocab = _initialize_vocabulary(special_tokens)
    if debug_on():
        print("[2/3] pre-tokenize", flush=True)
    pre_tokens = _pre_tokenize(input_path, special_tokens, max_bytes=max_bytes)
    if debug_on():
        print("[3/3] build BPE merges", flush=True)
    return _merge_tokens(pre_tokens, vocab_size, vocab)

処理の流れはコードの通り3ステップに分かれている。

train_bpe
├── _initialize_vocabulary()  初期語彙を作る
├── _pre_tokenize()           コーパスを事前トークンの頻度表へ変える
└── _merge_tokens()           頻出ペアを結合し、語彙表と結合規則を更新する

ステップ1: 初期語彙を作る (難易度: ez)

初期語彙は 0 - 255 の全バイト値に、<|endoftext|> などの特殊トークン (分割してはいけない区切り文字列) を加えたものだ。

work/assignment1-basics/cs336_basics/my_tokenizer.pyd8a9247

def _initialize_vocabulary(special_tokens: list[str]) -> dict[int, bytes]:
    """バイトレベル BPE 用の初期語彙を構築する。

    語彙は「番号 -> バイト列」の辞書である。最初は 1 バイト (0-255) を全部登録し、
    その後ろに特殊トークンを足す。どの言語の文字も最終的にはバイトの列なので、
    英語以外でも同じ初期語彙から始められる。

    Args:
        special_tokens (list[str]): 語彙に追加する特殊トークンの文字列リスト。
            例: ["<|endoftext|>"]。文書の区切りなど、途中で割ってはいけない特殊トークン。

    Returns:
        dict[int, bytes]: トークン ID からバイト列への写像。
            0-255 は各バイトそのもの、256 以降が特殊トークン。
    """
    vocab: dict[int, bytes] = {}
    # 1 バイトは 0-255 の 256 通り。まずは「文字」ではなく「バイト」を全部覚える。
    for i in range(256):
        # bytes([i]) は「値 i のバイトを 1 個」。bytes(i) は「0 が i 個」なので別物。
        vocab[i] = bytes([i])
    # 特殊トークンは人間が読める文字列なので、UTF-8 でバイト列に直してから登録する。
    # enumerate(xs, start=256) は (256, xs[0]), (257, xs[1]), ... と番号を 256 から振る。
    for i, token in enumerate(special_tokens, start=256):
        # str.encode(encoding) は str をその文字コードの bytes にする。
        vocab[i] = token.encode("utf-8")
    return vocab

たとえば special_tokens=["<|endoftext|>"] を指定した場合、語彙表は以下のようになる。

{
    0: b"\x00",            # バイト値 0
    1: b"\x01",            # バイト値 1
    # ...
    97: b"a",              # ASCII の a も 1 バイトとして入っている
    # ...
    255: b"\xff",          # バイト値 255。ここまでが初期 256 語彙
    256: b"<|endoftext|>", # 特殊トークンは 256 以降に足す
}

値は str ではなく bytes だ。文字 (str) を初期語彙にすると未見の文字が未知語になるが、bytes なら256個であらゆるUTF-8文字列を表現できる。実装も range(256) で済む。

ステップ2: コーパスを事前トークンの頻度表へ変える (難易度: ez)

low low lower という文字列を事前に分割せず、そのまま隣接バイトペアを数えると次のようになる。

{
    (b"l", b"o"): 3,  # low / low / lower の lo
    (b"o", b"w"): 3,  # 同上の ow
    (b"w", b" "): 2,  # 単語末尾と空白。事前トークン化しないとこれが残る
    (b" ", b"l"): 2,  # 空白と次単語の先頭
    (b"w", b"e"): 1,  # lower の we
    (b"e", b"r"): 1,  # lower の er
}

ここで問題になるのが (b"w", b" ") のような単語末尾と空白のペアだ。自然言語コーパスでは空白や句読点 (. や ! など) の出現頻度が非常に高いため、テキスト全体を無差別に結合していくと e や dog. のような単語境界を越えた汎用性の低いペアが語彙を圧迫してしまう。

これを防ぐのが正規表現を用いた事前トークン化 (Pre-tokenization) だ。単語や記号の境界であらかじめ区切り、境界をまたぐ不適切な結合を阻止する。

本実装では GPT-2 のパターンを整理し、'(?:[sdmt]|ll|ve|re) などの正規表現で分割している。

work/assignment1-basics/cs336_basics/my_tokenizer.pyd8a9247

# GPT-2 と同じ「BPE の前に、文を単語っぽい塊へ切る」正規表現。
# finditer は左から順に、下の枝を上から試して最初に成功したものを取る。
#
# 枝 (左が優先):
#   '(?:[sdmt]|ll|ve|re)     短縮形だけ。's 'd 'm 't 'll 've 're
#   ?\p{L}+                  任意の半角スペース 0 or 1 個 + 文字の連続
#                            \p{L} は Unicode の Letter。hello も こんにちは も当たる
#   ?\p{N}+                  任意の半角スペース 0 or 1 個 + 数字の連続 (\p{N} = Number)
#   ?[^\s\p{L}\p{N}]+        任意の半角スペース 0 or 1 個 + 空白でも文字でも数字でもないもの
#                            句読点や記号。, . ! ... など
#   \s+(?!\S)                この先に非空白が無い空白 = 行末の余り空白
#   \s+                      それ以外の空白の連続
#
# ポイント: 空白は独立した塊にせず、次の単語側につくことが多い。
# 例: 'Hello, I\'m 42.' のマッチ (引用符は文の中の 'm 用)
#   "Hello" | "," | " I" | "'m" | " 42" | "."
# " world" は 1 塊。"world" と " " には分かれない。
GPT2_PRETOKENIZE_PATTERN = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

たとえば Hello, I'm 42. をこの正規表現で分割すると次のようになる。

[
    "Hello",  # 文字の連続
    ",",      # 記号
    " I",     # 空白は次の単語側につく
    "'m",     # 短縮形
    " 42",    # 空白 + 数字
    ".",      # 記号
]

縮約表現 ('m) や記号 (, や .) が分離される一方、空白は単独トークンにならず後続の単語や数字の先頭に付与される ( I や 42)。

同じルールを low low lower に適用すると ["low", " low", " lower"] となり、単語境界をまたぐペア (b"w", b" ") は結合候補から完全に排除される。

train_bpe ではこの集計を _count_pretokens_from_text で行う。要点は特殊トークンの除外と正規表現によるカウントの2つだ。

まず特殊トークン (<|endoftext|> など) はステップ1ですでに語彙へ登録済みのため、通常の文字として分解・集計されないよう確実に除外する。

# 特殊トークンを正規表現の区切りにする。| でつなぐので複数個でも1回の走査で足りる。
# regex.escape は < や | を普通の文字として扱うため。空リストなら None。
delimiter = "|".join(map(regex.escape, special_tokens)) or None

if delimiter is None:
    # 区切りが無ければ全文を1文書として数える
    count_document(text, len(text))
else:
    last = 0
    # 特殊トークンの出現位置を左から順に探す
    for boundary in regex.finditer(delimiter, text):
        # [last, 特殊トークン直前) だけを数え、特殊トークン自体は飛ばす
        count_document(text[last : boundary.start()], boundary.end())
        last = boundary.end()  # 次は特殊トークンの直後から
    # 最後の特殊トークンより後ろの残り
    count_document(text[last:], len(text))

たとえば hello<|endoftext|>world の場合、<|endoftext|> をスキップして前後の通常テキスト (hello と world) だけを count_document に渡す。

各文書に対して正規表現を適用し、マッチした文字列を1バイトごとのタプルに変換して頻度を数える。

for match in regex.finditer(GPT2_PRETOKENIZE_PATTERN, document):
    # マッチした文字列を UTF-8 バイト列にする
    token_bytes = match.group(0).encode("utf-8")
    # 1 バイトずつタプルに分解して頻度を加算する
    # 例: b"low" -> (b"l", b"o", b"w")
    counts[tuple(bytes([byte]) for byte in token_bytes)] += 1

ここで得られるのはペアの頻度ではなく事前トークン (バイト列タプル) 自体の出現頻度表だ。low low lower なら次のようになる (ペア頻度の計算はステップ3で行う)。

{
    (b"l", b"o", b"w"): 1,                   # "low"
    (b" ", b"l", b"o", b"w"): 1,             # " low" (先頭空白つき)
    (b" ", b"l", b"o", b"w", b"e", b"r"): 1, # " lower"
}

ファイルを独立したチャンクへ切る

事前トークンのカウントはチャンクごとに独立しているため、ファイルをあらかじめ分割して並列処理できる。切れ目は文字数の等分ではなく特殊トークンの先頭へ寄せる。単語の途中で切ると同じ単語が2つのチャンクに分割され、正確な頻度を集計できなくなるためだ。

流れは次の3手順で、直列も並列も同じである。

_find_chunk_boundaries で特殊トークン先頭に寄せた境界を決める
_build_pre_tokenize_jobs で各チャンクの読み込み範囲 [start, end) をジョブにする
各ジョブを _count_file_range で数える

分かれるのは最後の実行方法だけだ。プロセスが1つなら for ループ、複数なら multiprocessing.Pool に載せる。

jobs = _build_pre_tokenize_jobs(input_path, special_tokens, num_chunks, max_bytes)
if num_processes == 1:
    return _run_pre_tokenize_serial(jobs)           # 1プロセスで順に数える
return _run_pre_tokenize_parallel(jobs, num_workers)  # Pool に載せる

どちらも戻り値はチャンクごとの頻度表のリストで、この時点では合算しない。足し算とペア結合はステップ3の逐次処理だ。全体の最頻ペアを1組ずつ選ぶ必要があるため、結合側は並列化していない。

並列処理側でワーカーが大きなチャンクの頻度表を保持したまま次のジョブへ進むと、メモリが解放されずに残り続けてしまう。maxtasksperchild=1 で1ジョブごとにプロセスを捨ててメモリをOSへ返す。

下のデモで特殊トークン先頭への寄せと文字数等分を比べる。

ステップ3: 頻出ペアを結合する

work/assignment1-basics/cs336_basics/my_tokenizer.pyd8a9247

def _merge_tokens(
    pre_tokens: list[dict[tuple[bytes, ...], int]], vocab_size: int, vocab: dict[int, bytes]
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """チャンクごとの頻度表をまとめてから、BPE マージする。

    先に各チャンクのカウントを 1 つの表に足す。そのあと、よく隣り合う 2 つの部品を
    1 つにまとめ、新しい語彙番号を付ける。

    Args:
        pre_tokens (list[dict[tuple[bytes, ...], int]]): チャンクごとの頻度表。
        vocab_size (int): 目標語彙サイズ (特殊トークンを含む)。
        vocab (dict[int, bytes]): 初期語彙。トークン ID からバイト列への写像。

    Returns:
        tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
            vocab: 訓練後の語彙。番号から部品 (バイト列) への辞書。
            merges: 作成順の BPE マージ。各要素は (左側の部品, 右側の部品)。
    """
    # ----- 引数が何かを先に固定する -----
    # pre_tokens:
    #   CPU ごとに数えた「単語っぽい塊 (pretok)」の回数表が、list で何枚か入っている。
    #   1 枚の例: { (b"l", b"o", b"w"): 5 }  は「low が 5 回」という意味。
    #   キーは隣り合う 2 つではない。単語全体の部品列。
    # vocab_size:
    #   終わったあとに vocab が持っていてよい部品の個数の上限。256 バイト + 特殊トークン + 結合で増えた分。
    # vocab:
    #   今知っている部品。キーが番号、値がその部品のバイト列。回数は無い。
    #   最初は 0 -> b"\x00" ... 255 -> b"\xff" と、特殊トークン。

    # dict(vocab): 呼び出し側の辞書を壊さないようコピーする。以降このコピーだけを育てる。
    vocab = dict(vocab)

    show = debug_on()
    merge_progress = my_utils._rich_progress() if show else None
    merge_task = None
    if merge_progress is not None:
        merge_progress.start()
        merge_task = merge_progress.add_task("combine counts", total=len(pre_tokens))
    try:
        pretok_freq: Counter[tuple[bytes, ...]] = Counter()
        for chunk_table in pre_tokens:
            pretok_freq.update(chunk_table)
            if merge_progress is not None and merge_task is not None:
                merge_progress.advance(merge_task)

        pair_freq: Counter[tuple[bytes, bytes]] = Counter()
        pair_index: dict[tuple[bytes, bytes], set[tuple[bytes, ...]]] = defaultdict(set)
        if merge_progress is not None and merge_task is not None:
            merge_progress.update(
                merge_task,
                description="count pairs",
                completed=0,
                total=len(pretok_freq),
            )
        for index, (parts, how_many) in enumerate(pretok_freq.items(), start=1):
            for pair in _adjacent_pairs(parts):
                pair_freq[pair] += how_many
                pair_index[pair].add(parts)
            if merge_progress is not None and merge_task is not None and index % 1000 == 0:
                merge_progress.update(merge_task, completed=index)

        merge_history: list[tuple[bytes, bytes]] = []
        remaining = max(vocab_size - len(vocab), 0)
        if merge_progress is not None and merge_task is not None:
            merge_progress.update(
                merge_task,
                description="bpe merge",
                completed=0,
                total=remaining,
            )

        while len(vocab) < vocab_size:
            if not pair_freq:
                break

            left, right = max(pair_freq, key=lambda pair: (pair_freq[pair], pair))
            merged_bytes = left + right

            # 該当 pretok を先に全部 pretok_freq から外す。
            # 同じマージで別 pretok が同じ新キーに着地しても、回数を混ぜてから pop しないため。
            affected = list(pair_index.get((left, right), ()))
            new_counts: Counter[tuple[bytes, ...]] = Counter()
            for parts in affected:
                how_many = pretok_freq.pop(parts, 0)
                if not how_many:
                    continue
                for pair in _adjacent_pairs(parts):
                    pair_freq[pair] -= how_many
                    if pair_freq[pair] <= 0:
                        del pair_freq[pair]
                    holders = pair_index[pair]
                    holders.discard(parts)
                    if not holders:
                        del pair_index[pair]
                new_counts[_replace_pair(parts, left, right, merged_bytes)] += how_many

            for new_parts, how_many in new_counts.items():
                pretok_freq[new_parts] += how_many
                for pair in _adjacent_pairs(new_parts):
                    pair_freq[pair] += how_many
                    pair_index[pair].add(new_parts)

            vocab[max(vocab) + 1] = merged_bytes
            merge_history.append((left, right))
            if merge_progress is not None and merge_task is not None:
                merge_progress.advance(merge_task)
    finally:
        if merge_progress is not None:
            merge_progress.stop()

    return vocab, merge_history

各チャンクの頻度表を合算して全体の pretok_freq を作成したらBPEの結合ループに入る。下のデモでは以下の4つのステップを順に進めて確認できる。

ペア頻度の集計: 各事前トークンの内部だけで隣接バイトペアを列挙し、出現回数を pair_freq に加算する (事前トークンの境界を越えたペアは作らない)。
最頻ペアの選出: 出現頻度が最も高いペアを選ぶ。同点の場合は決定論的な結果を得るためバイトペアの辞書順でタイブレークする (key=lambda pair: (pair_freq[pair], pair))。
貪欲な結合と記録: 選ばれたペアを事前トークン内で左から順に貪欲に結合し、新しいバイト列として語彙表 vocab に登録して結合したペアを結合規則 merges に追記する。テキストをエンコードする際に学習時と同じ順序を再現するため、merges の順序が必須となる。
目標サイズまで反復: この処理を目標の語彙サイズ (vocab_size) に達するまで繰り返す。

ペアを結合するたびにコーパス全体を走査すると遅すぎるため本実装では各ペアを含む事前トークンを逆引きできる pair_index を保持している。結合対象の事前トークンだけを更新し、消滅したペアと新しく生じたペアの頻度を差分更新することで大幅に高速化している。デモは差分更新ではなく手順のたびに数え直している。

encode と decode

train_bpe が vocab と merges を作る側であり、MyTokenizer はその規則を文字列と ID 列の相互変換に適用する側だ。

work/assignment1-basics/cs336_basics/my_tokenizer.pyd8a9247

class MyTokenizer:
    """訓練済みの語彙とマージ規則で、テキストとトークン ID を行き来する。

    train_bpe が規則を学習する側で、このクラスはその規則を適用する側である。
    encode は文字列を ID 列にし、decode は ID 列を文字列に戻す。
    """

    def __init__(
        self,
        vocab: dict[int, bytes],
        merges: list[tuple[bytes, bytes]],
        special_tokens: list[str] | None = None,
    ) -> None:
        """語彙・マージ・特殊トークンからトークナイザーを組み立てる。

        special_tokens が語彙に無いときは、語彙の末尾へ追加する想定である。
        encode では特殊トークン文字列を細切れにせず、常に 1 トークンとして扱う。

        Args:
            vocab (dict[int, bytes]): トークン ID からバイト列への写像。
                例: {0: b"h", 1: b"e", 2: b"he"}
            merges (list[tuple[bytes, bytes]]): BPE マージの作成順一覧。
                各要素 (token1, token2) は token1 と token2 を結合したことを表す。
                例: [(b"h", b"e")] は b"h" + b"e" -> b"he"。
            special_tokens (list[str] | None): 特殊トークン。None なら空と同じ。
                例: ["<|endoftext|>", "<|pad|>"]
        """
        # 呼び出し側の dict を壊さない。特殊トークン追加で書き込むため。
        self.vocab = dict(vocab)  # 例: {31373: b"hello"}
        self.merges = merges  # 例: [(b"h", b"e"), (b"he", b"llo")]
        self.special_tokens = list(special_tokens or [])  # 例: ["<|endoftext|>"]
        # regex の | は左から試すため、短い特殊トークンが長いものの一部に一致しないよう長い順にする。
        # 例: ["<|pad|>", "<|endoftext|>"] -> r"<\|endoftext\|>|<\|pad\|>"
        self._special_token_pattern = "|".join(
            regex.escape(special_token)
            for special_token in sorted(self.special_tokens, key=len, reverse=True)
        )

        # 語彙に無い特殊トークンは末尾 ID で足す (公式テストの前処理と同じ)。
        # 例: {0: b"a"} と ["<|pad|>"] なら {0: b"a", 1: b"<|pad|>"} にする。
        known_token_bytes = set(self.vocab.values())
        next_token_id = max(self.vocab, default=-1) + 1
        for special_token in self.special_tokens:
            # 例: "<|pad|>" -> b"<|pad|>"
            special_token_bytes = special_token.encode("utf-8")
            if special_token_bytes in known_token_bytes:
                continue
            self.vocab[next_token_id] = special_token_bytes
            known_token_bytes.add(special_token_bytes)
            next_token_id += 1

        # encode 用の逆引き。bytes -> ID、特殊トークン文字列 -> ID。
        # 例: {31373: b"hello"} -> {b"hello": 31373}
        self._bytes_to_id = {token_bytes: token_id for token_id, token_bytes in self.vocab.items()}
        # 例: {"<|pad|>": 50257}
        self._special_to_id = {
            special_token: self._bytes_to_id[special_token.encode("utf-8")]
            for special_token in self.special_tokens
        }
        # 作成が早い merge ほど rank が小さい。encode では小さい方を優先する。
        # 例: [(b"h", b"e"), (b"l", b"l")] -> {(b"h", b"e"): 0, (b"l", b"l"): 1}
        self._merge_ranks = {pair: rank for rank, pair in enumerate(self.merges)}

    @classmethod
    def from_files(
        cls,
        vocab_filepath: str | os.PathLike[str],
        merges_filepath: str | os.PathLike[str],
        special_tokens: list[str] | None = None,
    ) -> "MyTokenizer":
        """シリアライズした語彙とマージからトークナイザーを復元する。

        train_bpe_tinystories / train_bpe_owt が書く JSON と同じ形式を読む。
        vocab: {"0": [72], "1": [101], ...}  (キーは ID の文字列、値はバイト値の list)
        merges: [[[104], [101]], ...]         (各 merge は [左バイト列, 右バイト列])

        Args:
            vocab_filepath (str | os.PathLike[str]): 語彙ファイルのパス。
                例: "vocab.json"
            merges_filepath (str | os.PathLike[str]): マージ一覧ファイルのパス。
                例: "merges.json"
            special_tokens (list[str] | None): 特殊トークン。None なら空と同じ。

        Returns:
            MyTokenizer: 復元したトークナイザー。
        """
        with open(vocab_filepath, encoding="utf-8") as f:
            raw_vocab = json.load(f)
        with open(merges_filepath, encoding="utf-8") as f:
            raw_merges = json.load(f)

        # JSON のキーは常に str なので int に戻し、list[int] を bytes にする。
        vocab = {int(token_id): bytes(token_bytes) for token_id, token_bytes in raw_vocab.items()}
        merges = [(bytes(left), bytes(right)) for left, right in raw_merges]
        return cls(vocab, merges, special_tokens)

    def _split_with_special(self, text: str) -> list[str]:
        """特殊トークンを残したままテキストを区切る。

        重なり合う特殊トークンは長い方を優先する。
        例: "hello<|endoftext|>world"
            -> ["hello", "<|endoftext|>", "world"]
        """
        if not self.special_tokens:
            return [text] if text else []

        # 括弧付き split なので区切り自体もリストに残る。
        # 例: "hello<|endoftext|>world" -> ["hello", "<|endoftext|>", "world"]
        return [
            part
            for part in regex.split(f"({self._special_token_pattern})", text)
            if part
        ]

    def _apply_bpe(self, parts: list[bytes]) -> list[bytes]:
        """1 pre-token のバイト列に、rank がいちばん若い隣組を繰り返し結合する。

        例: [b"h", b"e", b"l", b"l", b"o"]
            -> [b"he", b"llo"]（対応する merges がある場合）
        """
        while len(parts) >= 2:
            # 今ある隣組のうち、訓練でいちばん早く作られた merge を選ぶ。
            # 例: (b"h", b"e") の rank 0 と (b"l", b"l") の rank 1 なら前者を選ぶ。
            best_i: int | None = None
            best_rank: int | None = None
            for i in range(len(parts) - 1):
                rank = self._merge_ranks.get((parts[i], parts[i + 1]))
                if rank is None:
                    continue
                if best_rank is None or rank < best_rank:
                    best_rank = rank
                    best_i = i
            # もう語彙にある merge が無いので終了。
            if best_i is None:
                break
            # 例: b"h" + b"e" -> b"he"
            merged = parts[best_i] + parts[best_i + 1]
            parts = parts[:best_i] + [merged] + parts[best_i + 2 :]
        return parts

    def _encode_ordinary(self, text: str) -> list[int]:
        """特殊トークンを含まない区間を pre-token -> BPE -> ID にする。

        例: "hello" -> [b"h", b"e", b"l", b"l", b"o"]
            -> [b"hello"] -> [31373]
        """
        ids: list[int] = []
        for match in regex.finditer(GPT2_PRETOKENIZE_PATTERN, text):
            # 訓練時と同じく、まず 1 バイトずつの列にする。
            parts = [bytes([byte]) for byte in match.group(0).encode("utf-8")]
            for token_bytes in self._apply_bpe(parts):
                ids.append(self._bytes_to_id[token_bytes])
        return ids

    def encode(self, text: str) -> list[int]:
        """入力文字列をトークン ID の列に符号化する。

        特殊トークンで区切ったあと、通常区間は pre-token に分け、
        各 pre-token に merges を適用して語彙 ID へ写す。

        Args:
            text (str): 符号化するテキスト。

        Returns:
            list[int]: トークン ID の列。

        Example:
            vocab に b"hello": 31373、特殊トークン: 50256 がある場合、
            "hello<|endoftext|>" -> [31373, 50256]
        """
        ids: list[int] = []
        for part in self._split_with_special(text):
            special_id = self._special_to_id.get(part)
            if special_id is not None:
                # 特殊トークンは切らず、語彙の 1 ID としてそのまま出す。
                ids.append(special_id)
            else:
                ids.extend(self._encode_ordinary(part))
        return ids

    def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
        """文字列のイテラブルを遅延的に符号化し、トークン ID を順に返す。

        大きなファイルを一度にメモリへ載せられないときに使う。
        例えばファイルハンドルを渡すと、チャンクごとに encode して ID を yield する。
        全部つなげた結果は、結合した文字列を encode した列と同じになる想定である。

        Args:
            iterable (Iterable[str]): 符号化する文字列の列。
                例: ["hello ", "world"]

        Yields:
            int: トークン ID。
        """
        # チャンクを結合しない。1 塊ずつ encode してすぐ yield する。
        # ファイルを行イテレートする使い方なら、行境界で pre-token もだいたい切れる。
        for chunk in iterable:
            yield from self.encode(chunk)

    def decode(self, ids: list[int]) -> str:
        """トークン ID の列を文字列に復号する。

        各 ID を語彙のバイト列に戻し、連結して UTF-8 でデコードする。

        Args:
            ids (list[int]): トークン ID の列。
                例: [31373]（vocab[31373] == b"hello" の場合）

        Returns:
            str: 復号したテキスト。上の例なら "hello"。
        """
        # 不正バイトが混ざっても落とさない。
        return b"".join(self.vocab[token_id] for token_id in ids).decode("utf-8", errors="replace")

encode は3手順で進む。

特殊トークンで区切る。重なるときは長い方を優先し <|endoftext|> のような文字列は細切れにせず 1 ID のまま出す。
残った文書は訓練と同じ GPT-2 の正規表現で事前トークン化し、1バイトずつの列にする。
その列の隣接ペアのうち結合規則 merges の作成順が一番早いもの (rank が小さいもの) を選んで結合する。これをもう使える結合規則が無くなるまで繰り返す。

def encode(self, text: str) -> list[int]:
    ids: list[int] = []
    # 括弧付き split なので区切りの特殊トークン自体もリストに残る。
    for part in self._split_with_special(text):
        special_id = self._special_to_id.get(part)
        if special_id is not None:
            ids.append(special_id)  # 特殊トークンは 1 ID
        else:
            ids.extend(self._encode_ordinary(part))  # 事前トークン化してから BPE
    return ids

結合の中身は今ある隣接ペアを rank で比べるだけだ。たとえば [b"h", b"e", b"l", b"l", b"o"] で (b"h", b"e") の rank が (b"l", b"l") より小さければ、先に b"he" ができる。

def _apply_bpe(self, parts: list[bytes]) -> list[bytes]:
    while len(parts) >= 2:
        best_i = None
        best_rank = None
        for i in range(len(parts) - 1):
            rank = self._merge_ranks.get((parts[i], parts[i + 1]))
            if rank is None:
                continue  # 訓練で作っていない隣接ペアは結合しない
            if best_rank is None or rank < best_rank:
                best_rank = rank
                best_i = i  # 作成が早い結合規則を優先する
        if best_i is None:
            break
        merged = parts[best_i] + parts[best_i + 1]
        parts = parts[:best_i] + [merged] + parts[best_i + 2 :]
    return parts

大きなファイルは encode_iterable に渡す。渡した文字列をメモリ上でつなげず1件ずつ encode して ID を返す。行ごとに渡すなら事前トークンもだいたい行で切れるが、文字列の途中で単語が割れると全文を一度に encode した結果と一致しなくなる。

decode は逆向きで ID を語彙表のバイト列に戻して連結し、UTF-8 として読む。途中に不正なバイトがあってもそこで止めず、その箇所は置換文字にする。

def decode(self, ids: list[int]) -> str:
    # 不正バイトが混ざっても落とさない。
    return b"".join(self.vocab[token_id] for token_id in ids).decode("utf-8", errors="replace")

下のデモは hello<|endoftext|>low を encode の3手順から decode まで1手ずつ進める。結合規則はこの例に固定してある。

訓練結果は JSON に書いておき from_files でこのクラスへ戻す。次節の encode の速度は、こうして復元したトークナイザーをコーパスにかけたときの値だ。

実装の検証とコーパス訓練

TinyStories と OpenWebText の大規模コーパスを用いて実際にトークナイザーを訓練した計測値が以下だ。

訓練時間・メモリ・最長トークン

コーパス	語彙サイズ	訓練時間	ピーク RSS	最長トークン
TinyStories	10,000	約 136 秒	約 0.14 GiB	先頭スペース付き `accomplishment` (15 バイト)
OpenWebText	32,000	約 16,559 秒 (約 4.6 時間)	約 8.81 GiB	64 バイトの `ÃÂ` (`\xc3\x83\xc3\x82`) の繰り返し

TinyStories の最長トークンは子供向け英語の頻出語 ( accomplishment) だったのに対し、OpenWebText は Web クロール由来の文字化け (UTF-8 の多重エンコード残骸) だった。

ボトルネック

プロファイラの計測では語彙サイズ 10,000 の TinyStories では事前トークン化が最長だった。一方語彙サイズ 32,000 の OpenWebText では最頻ペア選択 (pick_best_pair) が約 90% を占め、結合処理側がボトルネックになっていた。

圧縮率 (各コーパスから 10 文書)

テキスト	tokenizer	bytes	tokens	bytes/token
TinyStories	TinyStories (vocab 10,000)	10,909	2,701	4.04
OpenWebText	OpenWebText (vocab 32,000)	50,784	11,245	4.52
OpenWebText	TinyStories (vocab 10,000)	50,784	14,887	3.41

同じ OpenWebText 10 文書を TinyStories 用トークナイザーで処理すると、専用の OWT トークナイザーに比べてトークン数が約 1.32 倍に増える。バイトレベル BPE のため未知語にはならないが、コーパスのドメインがずれると分割が細かくなり圧縮効率が落ちる様子が確認できる。

encode のスループット

約 1.49 MB/s (1.49×10^6 bytes/sec) だった。この速度で The Pile (825 GiB) を単一プロセスでエンコードすると約 165 時間 (約 6.9 日) かかる計算になる。

uint16 への export

corpus	split	tokens	max_id	elapsed
TinyStories	valid	5,461,210	9,999	約 20.5 秒
TinyStories	train	540,796,778	9,999	約 2,223 秒 (約 37 分)
OpenWebText	valid	66,401,098	31,999	約 366 秒 (約 6 分)
OpenWebText	train	2,727,120,452	31,999	約 10,814 秒 (約 3.0 時間)

どちらも語彙サイズが 65,535 以下のためトークン ID 列は uint16 でコンパクトに保存できる。

おわりに

今回のように各セクションで学んだことをまとめて記事にしていこうと思う。

次回: Assignment 1 Section 3: Transformer Language Model Architecture

偏向報道 ↩
Section 1はUTF-8やPythonの基礎環境に関する内容なので割愛した。今後も準備セクションは適宜スキップする。 ↩
UTF-8ならASCIIを1バイトで表せ、バイトオーダーも不要。英語中心のコーパスをコンパクトに扱いつつ、任意のUnicode文字を過不足なく表現できる。 ↩

barabara diary

Stanford CS336 Assignment 1 Section 2: BPEトークナイザー

インドア系ならトランスフォーマー

Stanford CS336

トークナイザーとは何か

`train_bpe`の全体像

ステップ1: 初期語彙を作る (難易度: ez)

ステップ2: コーパスを事前トークンの頻度表へ変える (難易度: ez)

ファイルを独立したチャンクへ切る

ステップ3: 頻出ペアを結合する

encode と decode

実装の検証とコーパス訓練

訓練時間・メモリ・最長トークン

ボトルネック

圧縮率 (各コーパスから 10 文書)

encode のスループット

uint16 への export

おわりに

事前トークン

pair_freq

merges

vocab

トークン列

結合規則

decode

バイト語彙 (ID 0〜255)

文字語彙との対比

特殊トークンの先頭で切る

文字数で等分する

AIレビュー

インドア系ならトランスフォーマー

Stanford CS336

トークナイザーとは何か

train_bpeの全体像

ステップ1: 初期語彙を作る (難易度: ez)

ステップ2: コーパスを事前トークンの頻度表へ変える (難易度: ez)

ファイルを独立したチャンクへ切る

ステップ3: 頻出ペアを結合する

encode と decode

実装の検証とコーパス訓練

訓練時間・メモリ・最長トークン

ボトルネック

圧縮率 (各コーパスから 10 文書)

encode のスループット

uint16 への export

おわりに

Footnotes

事前トークン

pair_freq

merges

vocab

トークン列

結合規則

decode

バイト語彙 (ID 0〜255)

文字語彙との対比

特殊トークンの先頭で切る

文字数で等分する

AIレビュー

`train_bpe`の全体像