0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@Snowlan(川口 海龍)

ゴミから最小のBPE実装へ #14

0
Posted at

本当にやるの？

正直、Pythonでの実装方法なんて、AIに聞いた方が圧倒的に速い。

AIなら、おそらく数秒できれいなコードを書いてくれるだろう。しかし、川口はPythonの知識がほぼゼロなので、8時間以上かかってしまった。

川口みたいに、どうしても自分で実装してみたいという方は、Pythonの書き方をネットの記事で調べながら進めてもいいだろう。

もちろん、AIに聞いてもいい。ただし、完成したコードを教えてもらうのではなく、Pythonの文法に関するヒントだけをもらうように指示しておこう。

最初はゴミコードでもいいから、とにかく自分で書いてみる。そして、少しずつバージョンアップしていこう。

最後は、夜中にPCを起動し、ヘッドホンから流れる音楽に力をもらいながら、独り言をつぶやきつつ、ゴールに向かって進もう！

ここはゴミ箱ではないので、川口が最初に書いたゴミコードはここには載せたくない。

恥ずかしいが、GitHubにはアップロードしておこうと思う。

川口のゴミコードを見て、「この程度のゴミコードなら、自分にも書けそうだ」と思ってもらえたら、それで十分だ。

さあ、一緒にゴミを作ろう！

やりたいのは何？

今回実装したい機能は、以下の4つ。

  1. 対象の文字列から出現頻度の高いペアを調べる。
  2. [1.]で確認できたペアを、新しいTokenに置換する。
  3. 置換後のToken列に対して、[1.]から[2.]を繰り返しす。
  4. 置換したToken列から、元の文字列を復元する。

Wikipediaのバイト対符号化によると、decode処理がないようだが、やってみようと、正しくできているのか分からないので、decodeもやってみよう。
頭の中にpythonが実行できる方は↑を無視してください。ご自由に

入力データを準備する

今回もWikipediaのバイト対符号化の例の文字列を使用する。

input = "ABCDCDABCDCDE"
tokens = list(input.encode("utf8"))

print(tokens)

実行結果：

[65, 66, 67, 68, 67, 68, 65, 66, 67, 68, 67, 68, 69]

また、Byteの値は0～255なので、新しいToken IDは256から割り当てることにした。

next_token_id = 256
pair_size = 2

BPE関数を実装する

BPEの処理を大きく分けると、2つになる。

  1. 最も出現頻度の高いペアを調べる。
  2. そのペアを新しいToken IDに置換する。

出現頻度の高いペアを調べる

まず、となり2つのTokenを取り出し、出現回数を数える。

counts = {}

for index in range(tokens_from_utf8_len - 1):
    pairs = tokens_from_utf8[index:index + pair_size]
    counts[tuple(pairs)] = counts.get(tuple(pairs), 0) + 1

max_pair = max(counts, key=counts.get)

ここでは、Pythonの辞書を使って、ペアと出現回数の関係を保存する。

例えば、

(67, 68) : 4
(65, 66) : 2

のような形になる。

counts.get(tuple(pairs), 0) + 1は、すでに存在するペアなら出現回数を1増やし、存在しなければ1から数える。

なお、辞書のKeyにはListを使用できないため、tuple()に変換している。

最後に、

max_pair = max(counts, key=counts.get)

で、最も出現回数の多いペアを取得する。

今回の入力データの場合、最初に選ばれるペアは(67, 68)、つまりCDになる。

ここでは出現回数が同じペアが複数存在した場合の処理は、面倒くさいので特別に実装していない。
ちなみに、今回のコードでは、出現回数が同じ場合、辞書に先に登録されたペアが選ばれる。意図的に設計したわけではなく、Pythonの仕様に任せているだけだ。

ペアを新しいTokenに置換する

次に、先ほど取得したmax_pairを新しいToken IDに置換する。

while i < tokens_from_utf8_len:

    if tuple(tokens_from_utf8[i:i + pair_size]) == max_pair:
        new_tokens.append(next_token_id)
        i += pair_size
    else:
        new_tokens.append(tokens_from_utf8[i])
        i += 1

ここで重要なのは、iの増加量だ。

置換対象のペアを発見した場合、2つのTokenを1つにまとめるため、i += 2で次に進む。

対象のペアがなければ、元のTokenをそのまま追加して、i += 1で進む。

つまり、ペアを数える処理ではとなり位置をすべて確認するが、置換する処理では、すでに置換したTokenを再利用しない。

前回の記事で考えた、ペアの重複問題に対応するための処理でもある。

BPE関数の全体

ここまでの処理を関数にまとめる。

def BPE(tokens_from_utf8, next_token_id):
    new_tokens = []
    tokens_from_utf8_len = len(tokens_from_utf8)
    counts = {}
    i = 0

    for index in range(tokens_from_utf8_len - 1):
        pairs = tokens_from_utf8[index:index + pair_size]
        counts[tuple(pairs)] = counts.get(tuple(pairs), 0) + 1

    max_pair = max(counts, key=counts.get)

    while i < tokens_from_utf8_len:
        if tuple(tokens_from_utf8[i:i + pair_size]) == max_pair:
            new_tokens.append(next_token_id)
            i += pair_size
        else:
            new_tokens.append(tokens_from_utf8[i])
            i += 1

    return new_tokens, next_token_id, max_pair

戻り値は3つにした。

  • new_tokens：置換後のToken列
  • next_token_id：今回割り当てたToken ID
  • max_pair：置換対象となったペア

最初は置換後のToken列だけあればよいと思ったが、後でDecodeするためには、どのToken IDがどのペアに対応しているかも必要になる。

そのため、max_pairとnext_token_idも戻すことにした。

BPEを繰り返す

ここまでで、1回分のBPE処理は実装できた。

ただし、BPEは1回置換して終了するわけではない。

置換した結果に対して、再び出現頻度の高いペアを探し、新しいTokenに置換する必要がある。

そこで、再帰関数を使って繰り返すことにした。

tokens_table = {}

def loop_BPE(tokens, next_token_id, n):

    if n == 0:
        return tokens, tokens_table

    else:
        new_token, token_ids, max_pair = BPE(tokens, next_token_id)

        tokens_table[token_ids] = max_pair

        next_token_id += 1

        return loop_BPE(new_token, next_token_id, n - 1)

今回は実験なので、繰り返す回数をnで指定する。

new_token, tokens_table = loop_BPE(tokens, next_token_id, 3)

これでBPEを3回実行する。

ここで、もう1つ必要なものがある。

tokens_table[token_ids] = max_pair

これは、新しく作ったToken IDと、置換前のペアの関係を保存するための処理だ。

今回の実行結果は、以下のようになる。

256 : (67, 68)
257 : (65, 66)
258 : (257, 256)

注目したいのは、258の中身だ。

258 -> (257, 256)

この257と256も、元のByteではなく、BPEによって新しく作られたTokenだ。

つまり、BPEを繰り返すと、新しいTokenの中に、さらに別の新しいTokenが含まれるようになる。

この関係を保存しておけば、後で元のByte列に戻せるはずだ。

実行結果を確認する

今回、3回のBPE処理を実行した。

1回目

最初の文字列は、

ABCDCDABCDCDE

13個のTokenになる。

ここで最も出現回数の多いCDを、Token ID 256に置き換える。

256 -> (67, 68)

結果：

[65, 66, 256, 256, 65, 66, 256, 256, 69]

13個から9個になった。

2回目

次に、置換後のToken列を対象に、再び出現頻度の高いペアを調べる。

今回は(65, 66)が選ばれた。

257 -> (65, 66)

結果：

[257, 256, 256, 257, 256, 256, 69]

9個から7個になった。

3回目

さらに同じ処理を繰り返す。

今回は(257, 256)が選ばれた。

258 -> (257, 256)

結果：

[258, 256, 258, 256, 69]

7個から5個になった。

最初の13個のTokenが、3回のBPE処理によって5個になった。

ただし、これはToken列が短くなったという意味であり、保存するデータ全体のサイズが小さくなったわけではない。

新しく作ったToken IDとペアの関係も、別途保存する必要があるためだ。

Decodeを実装する

ここまでは、元のToken列を短くする処理だった。

次に、この5個のTokenから、元の文字列を復元できるか確認する。

最初は、tokens_tableに保存したペアを使えば、簡単に戻せると思った。

しかし、実際には少し問題がある。

例えば、

258 -> (257, 256)
257 -> (65, 66)
256 -> (67, 68)

という関係がある場合、258を展開しても、すぐに元のByte列には戻らない。

258
 ↓
257 256
 ↓
65 66 67 68

のように、複数回の展開が必要になる。

そこで、今回は新しく作ったToken IDの大きい順に、ペアへ置き換えていくことにした。

def decode(tokens, tokens_table):
    new_tokens = tokens.copy()

    for k in sorted(tokens_table.keys(), reverse=True):
        result = []

        for token in new_tokens:
            if token == k:
                result.extend(tokens_table[k])
            else:
                result.append(token)

        new_tokens = result

    return new_tokens

ここで重要なのは、

sorted(tokens_table.keys(), reverse=True)

の部分だ。

今回の実装では、BPEの処理が進むたびに新しいToken IDを1ずつ増やしている。

そのため、後から作ったTokenは、先に作ったTokenを含む場合があるが、その逆はない。

新しいTokenから古いTokenへ順番に展開すれば、最終的には元のByte列に戻せる。

また、

result.extend(tokens_table[k])

を使用しているのは、1つのTokenを2つのTokenへ展開するためだ。

append()ではペア全体がListとして追加されてしまうため、今回はextend()を使用した。

元の文字列に戻せるか？

最後に、Decodeした結果をUTF-8で文字列へ戻す。

after_decode = decode(new_token, tokens_table)

output = bytes(after_decode).decode("utf-8")

print("input : ", input)
print("output : ", output)

print("old_tokens = new_tokens ? ", tokens == after_decode)
print("input = output ? ", input == output)

実行結果：

input :  ABCDCDABCDCDE
output :  ABCDCDABCDCDE

old_tokens = new_tokens ?  True
input = output ?  True

これで、今回の入力データについては、BPEで置換したToken列から元の文字列を復元できることが確認できた。

今回の実装では、

13 Tokens
    ↓
 9 Tokens
    ↓
 7 Tokens
    ↓
 5 Tokens
    ↓ Decode
13 Byte Tokens
    ↓ UTF-8
元の文字列

という一連の処理を再現できた。

まだ実装していないこと

今回のコードは、あくまでBPEの基本的な動作を確認するための最小実装だ。

例えば、以下の問題には対応していない。

  • 出現回数が同じペアが複数存在した場合の明示的な優先順位
  • 入力Tokenが0個または1個の場合の処理
  • どの条件でBPEを終了するかの判定
  • 学習したMerge Ruleを保存し、別の文字列へ適用する処理
  • Vocabulary Sizeの制限

特に、今回実装したのは、入力データから頻出ペアを調べて新しいTokenを作る処理だ。

実際のTokenizerでは、学習時に作成したVocabularyやMerge Ruleを使って、別の入力データをEncodeする必要がある。

そのため、今回のコードをそのまま実用的なTokenizerとして使用することはできない。

まとめ

今回、Pythonを使ってBPEの最小限の機能を実装してみた。

実際にコードを書いたことで、BPEの処理を以下のように分けて考えられるようになった。

  1. 出現頻度の高いペアを探す。
  2. ペアに新しいToken IDを割り当てる。
  3. 元のToken列を新しいTokenに置換する。
  4. Token IDとペアの関係を保存する。
  5. 必要に応じて繰り返す。
  6. 保存した関係を逆にたどって元のByte列に戻す。

最初は、BPEとは頻出するペアを置換するだけのアルゴリズムだと思っていた。

しかし、実装してみると、ペアを探すこと自体よりも、新しいTokenをどう管理し、どう繰り返し、どう元に戻すかを考える必要があった。

また、今回は3回の処理を指定しただけだが、実際のTokenizerでは、どのペアを選び、どこで処理を終了するかによってVocabularyが変わる。

結局、BPEの基本的な処理はシンプルでも、Tokenizerとして利用するためには、まだ考えることが多そうだ。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?