トークナイザを手で作る

目標

バイトレベルのBPEを、標準ライブラリだけで自分で作ります。UTF-8のバイトから出発して隣り合うペアを数え、最もよく現れるペアをマージするルールを学習し、そのルールを学習した順序で適用してエンコードしたあと、元の文章に正確に戻します。最後に、語彙サイズを変えながらトークン数が減っていく曲線を測り、同じ内容の韓国語の段落と英語の段落を1つの語彙でエンコードして、トークン数を並べて見ます。

なぜ重要なのか

料金もコンテキストの上限も単位はトークンですが、トークンは文字でも単語でもありません。どの断片を1つと数えるかは、語彙を作るときに決まり、語彙はモデルごとに違います。そのため、「文字数に何倍かを掛ければいい」という見積もりは、言語が変わった瞬間に崩れます。ハングルの音節1つはUTF-8で3バイトなので、語彙が小さいと、1文字がトークン3つになります。このラボは、実際のモデルのトークナイザーを呼びません。このPodにはtransformers・tokenizers・tiktokenがなく、numpyもシステムのPythonにはありません。代わりに、同じアルゴリズムを手で作り、自分で作った語彙で測った数字だけを使います。そのため、「あるモデルはこの文を何トークンと数える」といった話は、ここではしません。アルゴリズム自体は、1文です。最もよく隣り合って現れる2つを1つにマージする作業を、語彙が望むサイズになるまで繰り返します。難しいのは細部です。重なる位置をどう数えるか、同点をどう解消するか、エンコードするときにルールをどの順序で適用するかが決まっていなければ、同じ文章で動かしても、語彙が毎回変わります。採点ツールは、書いておいた説明を信用しません。自分で作ったモジュールを実際に呼び出し、毎回異なる入力で関数を直接叩いて、採点ツールが別に計算した値と照合します。入力は実行のたびに変わるので、値を暗記して入れることはできません。

ステップ

/root/work/tf-token/bpe.pyに、サンプルの段落SAMPLE_KO・SAMPLE_ENと、to_ids(text)・from_ids(ids)を作成してください。文章をUTF-8のバイトのリストに開き、再び文章に戻します。
count_pairs(ids)を追加して、隣り合う2つが何回続けて現れるかを数えるようにしてください。キーは(앞, 뒤)(プレースホルダーは前と後ろです)のペアで、重なる位置もそのまま数えます。
merge(ids, pair, new_id)を追加して、そのペアが現れる位置を新しい番号1つに置き換えるようにしてください。左から、重ならないように進みます。
MIN_PAIR_COUNT = 2とtrain(text, vocab_size)を作成して、マージルールを学習した順序で集めるようにしてください。新しい番号は256から1つずつ増えていきます。
encode(text, merges)を作成して、学習した順序のままマージを適用するようにしてください。
decode(ids, merges)を作成して、番号を最後まで展開し、元の文章に戻すようにしてください。語彙になかった文字が混ざっていても、元の文章と1文字も違ってはいけません。
vocab_curve(text, sizes)を作成して、語彙サイズごとに同じ文章が何トークンになるかを測るようにしてください。返す値は、(어휘크기, 토큰수)(プレースホルダーは語彙サイズとトークン数です)のペアのリストです。
語彙を2つ学習して同じ2つの段落を測り、結果を、JSONのレポート(/root/work/tf-token/token_report.json)と、Markdownのレポート(/root/work/tf-token/token_report.md)に記録してください。

参考

実行の契約: 採点ツールは、/root/work/tf-token/bpe.pyをPythonモジュールとして読み込み、SAMPLE_KO・SAMPLE_EN・to_ids・from_ids・count_pairs・merge・MIN_PAIR_COUNT・train・encode・decode・vocab_curveを直接使います。スクリプトとして実行しないので、if __name__ == "__main__"はなくてかまいません。
サンプルの段落は、同じ内容を2つの言語で書きます。SAMPLE_KOはハングルが大半の300文字以上、SAMPLE_ENはASCIIだけを使う300文字以上である必要があります。内容は自由に決めてください。
to_ids(text)は、text.encode("utf-8")のバイトを整数のリストにして返します。文字の番号(ord)ではありません。
from_ids(ids)は、途切れたバイトが入ってきても、例外を出してはいけません。decodeのerrors引数を使ってください。
count_pairs([9, 9, 9])は{(9, 9): 2}です。重なる位置をそのまま数えます。
merge([5, 5, 5], (5, 5), 300)は[300, 5]です。前の2つを飲み込んだあとに残った5は、相手を失います。受け取ったリストをその場で書き換えず、新しいリストを作ってください。
trainの1周: 隣接するペアを数え、最も多く現れたペアを選び、そのペアを新しい番号にマージします。同点なら、(앞, 뒤)(プレースホルダーは前と後ろです)が辞書順で小さいペアを選びます。選んだペアがMIN_PAIR_COUNTより少ない回数しか現れなければ、止まります。返す値は、((앞, 뒤), 새번호)(プレースホルダーは前、後ろ、新しい番号です)のペアのリストで、順序がそのまま学習した順序です。
train(text, 256)は空のリストです。0から255はすでにバイトが使っていて、新しい番号を作る場所がありません。
encodeは、mergesを受け取った順序のまま、1回ずつ適用します。並べ替えたり、繰り返したりしないでください。
decodeは、新しい番号を2つに展開し、その2つの中にまた新しい番号があれば、さらに展開します。残るものがなくなったら、from_idsで読みます。
ステップ8のレポートでは、語彙を2つ学習します。どちらもサイズは512です。1つは、SAMPLE_KOとSAMPLE_ENを改行1つでつないだ文章で学習した共用語彙で、もう1つは、SAMPLE_ENだけで学習した英語専用語彙です。2つの語彙で同じ2つの段落をそれぞれエンコードして、トークン数を比べます。
vocab_sizeは上限にすぎません。マージする価値のあるペアがなくなれば、MIN_PAIR_COUNTで止まるので、実際のルールの数は、512から256を引いた値より少ないことがあります。そのため、shared_rules・en_only_rulesには、実際に学習したルールの数を書きます。
曲線は、つないだ文章について、サイズ256・320・384・512で測ります。
このPodにはインターネットがありません。pip installはできず、transformers・tokenizers・tiktokenもありません。numpyは/opt/onnx-lab/bin/pythonの中にしかないので、システムのPythonではimport numpyができません。標準ライブラリだけで十分です。
公式ドキュメント: BPEの原論文・Hugging Face — Byte-Pair Encoding tokenization・Attention Is All You Need
よくある間違い: バイトの代わりに文字番号を使う、ペアを2つずつ飛ばして数える、マージするときに1つずつしか進まずに重ねて飲み込む、同点のルールを決めない、エンコードするときにルールを並べ替える、デコードを1段階だけ展開して終わる。

テキストをバイトに分解する

/root/work/tf-token/bpe.pyに、サンプルの段落SAMPLE_KO(ハングル中心の300文字以上)・SAMPLE_EN(ASCIIだけの300文字以上)と、to_ids(text)・from_ids(ids)を作成してください。to_idsはUTF-8のバイトを整数のリストにして返し、from_idsはそれを再び文章に戻します。

text.encode("utf-8")がバイト列を返し、list()で包むと、0から255までの整数のリストになります。戻すときは、bytes(ids)でまとめてdecodeします。文字の途中で切れた断片が入ってきても例外を出してはいけないので、errors引数を渡してください。2つのサンプルは、同じ内容を2つの言語で書かなければ、あとで比較できません。

隣り合う2つを数える

count_pairs(ids)を追加してください。隣り合う2つが何回続けて現れるかを数えて、{(앞, 뒤): 횟수}(プレースホルダーは前、後ろ、回数です)を返します。重なる位置もそのまま数えるので、count_pairs([9, 9, 9])は{(9, 9): 2}です。

zip(ids, ids[1:])で、隣り合うペアを一度に走査できます。2つずつ飛ばして数えると、重なる位置を見落とし、そうなると、何をマージすれば得かの判断がずれます。リストが空か、要素が1つだけなら、空の辞書です。

1つのペアを1つにマージする

merge(ids, pair, new_id)を追加してください。pairが現れる位置を、新しい番号1つに置き換えます。左から、重ならないように進むので、merge([5, 5, 5], (5, 5), 300)は[300, 5]です。受け取ったリストはそのままにして、新しいリストを返してください。

インデックスを手で動かすwhile文が、最も正確です。ペアを見つけたら2つ進み、そうでなければ1つだけ進みます。1つずつしか進まないと、いま作った新しい番号を、再びペアの前の側として見てしまい、重ねて飲み込みます。リストをその場で書き換えると、呼び出した側が持っていた値が、黙って変わってしまいます。

マージルールを学習する

MIN_PAIR_COUNT = 2とtrain(text, vocab_size)を作成してください。隣接するペアを数え、最も多く現れたペアを選び(同点なら、(앞, 뒤)(プレースホルダーは前と後ろです)が辞書順で小さいペア)、そのペアを新しい番号にマージする作業を繰り返します。新しい番号は256から1つずつ増え、返す値は、((앞, 뒤), 새번호)(プレースホルダーは前、後ろ、新しい番号です)のペアのリストです。

1周ごとに数え直す必要があります。マージすると、隣り合う関係が変わるからです。max(counts.items(), key=...)のkeyで、同点のルールまで一度に書けます。選んだペアがMIN_PAIR_COUNTより少ない回数しか現れなければ、止まってください。1回しか現れないペアをマージしても、語彙が増えるだけで、トークンは1つも減りません。vocab_sizeが256なら、新しい番号を作る場所がないので、空のリストです。

学習した順序でエンコードする

encode(text, merges)を作成してください。文章をバイトに開いたあと、mergesを受け取った順序のまま、1回ずつ適用します。並べ替えたり、繰り返したりしないでください。

あとで学習したルールは、前に作った番号を材料として使います。そのため、順序を変えると、同じルールのリストでも違う結果が出ます。関数は3行で終わります。バイトに開き、ルールごとにmergeを呼び、残ったリストを返します。

1文字も違わずに元に戻す

decode(ids, merges)を作成してください。新しい番号を2つに展開し、その2つの中にまた新しい番号があれば、さらに展開します。残るものがなければバイトだけが残るので、それを文章として読みます。学習になかった文字が混ざった文章も、元の文章と1文字も違ってはいけません。

{새번호: (앞, 뒤)}(プレースホルダーは新しい番号、前、後ろです)の表を作っておけば、展開する作業は、その表を見るだけです。1段階だけ展開して終わってはいけません。展開した値が、また新しい番号であることがあります。すべて展開したあとは、from_idsがそのまま使えます。語彙になかった文字も、バイトとしては必ず書かれているので、この構造では、初めて見る文字が来ても失敗しません。

語彙を大きくするとトークンが減る

vocab_curve(text, sizes)を作成してください。sizesの語彙サイズごとに、その文章で最初からルールを学習し、同じ文章をエンコードして、トークン数を測ります。返す値は、(어휘크기, 토큰수)(プレースホルダーは語彙サイズとトークン数です)のペアのリストで、順序はsizesと同じです。

サイズごとに、そのサイズのルールで測る必要があります。最大の語彙で一度学習しておいて、すべての枠をそのルールで測ると、曲線が平らになります。前に作ったtrainとencodeをそのまま使えば、関数は5行です。参考までに、このアルゴリズムは前だけを見て選ぶので、train(글, 300)の結果は、train(글, 400)の結果の前半と同じです(プレースホルダーはテキストです)。サイズを大きくするのは、より長く動かすことであって、選び直すことではありません。減る幅がどう変わるかを、目で見てください。

語彙は誰の文章から作られたのか

語彙を2つ学習してください。1つは、SAMPLE_KOとSAMPLE_ENを改行1つでつないだ文章で学習した、サイズ512の共用語彙、もう1つは、SAMPLE_ENだけで学習した、サイズ512の英語専用語彙です。2つの語彙で同じ2つの段落をそれぞれエンコードしたあと、/root/work/tf-token/token_report.jsonに、vocab_size・ko_chars・ko_bytes・en_chars・en_bytes・shared_rules・en_only_rules・ko_tokens_shared・en_tokens_shared・ko_tokens_en_only・en_tokens_en_only・ko_bytes_per_token・en_bytes_per_token・curve・roundtrip_okを、/root/work/tf-token/token_report.mdに、## 무엇을 쟀나 ## 어휘를 키우면 토큰이 어떻게 줄었나 ## 한국어가 손해를 보는 이유 ## 바이트 수준이라 안 깨지는 것(韓国語の見出しで、順に「何を測ったか」「語彙を大きくするとトークンがどう減ったか」「韓国語が不利になる理由」「バイトレベルなので壊れないもの」を意味します)の4つの節で書いてください。

数字は手で書かずに、自分のコードを実際に動かして得た値で埋めてください。shared_rules・en_only_rulesは、train()が実際に返したルールの数です。サイズ512を渡しても、マージするペアがなくなれば、その前に止まります。ko_bytes_per_token・en_bytes_per_tokenは、共用語彙を基準に、バイト数をトークン数で割った値です。curveは、つないだ文章について、サイズ256・320・384・512で測った[[크기, 토큰수], ...](プレースホルダーはサイズとトークン数です)です。roundtrip_okは、2つの語彙の両方で、2つの段落が元の文章に正確に戻るかどうかです。英語専用語彙で韓国語をエンコードしても元に戻るというのが、バイトレベルの性質です。