tiktoken の decode_tokens_bytes でトークン分割されたバイト文字列を、分割されたまま復元する【日本語対応】

GPT の API を使うときに気になるトークン数はtiktokenで計算できる。tiktoken でトークン数を取得するのは簡単なのだが、トークン分割された後の文字列を取得するのは自前の実装が必要。今回はこのトークン分割されたバイト配列から、分割されたまま文字列を復元する方法について紹介する。

トークン分割されたバイト配列を取得する

例えば"こんにちわ"をトークン分割してバイト配列にするには下記。

main.py

text = "こんにちわ"
# gpt-4-turboはcl100k_baseを指定
enc = tiktoken.get_encoding("cl100k_base")
tokens = enc.encode(text)
bs = enc.decode_tokens_bytes(tokens)
print(bs)

結果としては以下のようなバイト配列となる。

[b'\xe3\x81\x93\xe3\x82\x93\xe3\x81\xab', b'\xe3\x81\xa1', b'\xe3\x82\x8f']

'\xe3\x81\x93\xe3\x82\x93\xe3\x81\xab'が「こんに」、'\xe3\x81\xa1'が「ち」、'\xe3\x82\x8f'が「わ」だ。復元はdecode()を使えばいける

bs = enc.decode_tokens_bytes(tokens) # [b'\xe3\x81\x93\xe3\x82\x93\xe3\x81\xab', b'\xe3\x81\xa1', b'\xe3\x82\x8f']
print(bs[0].decode()) # こんに
print(bs[1].decode()) # ち
print(bs[2].decode()) # わ

では"こんに夜"の場合はどうなるか。バイト配列は下記になる。

[b'\xe3\x81\x93\xe3\x82\x93\xe3\x81\xab', b'\xe5\xa4', b'\x9c']

最初の'\xe3\x81\x93\xe3\x82\x93\xe3\x81\xab'は同じく「こんに」だ。
では次の'\xe5\xa4'をデコードするとどうなるか、エラーになる。

UnicodeDecodeError: 'utf-8' codec can't decode bytes in position 0-1: unexpected end of data

これは'\xe5\xa4'単体ではデコードできず、次のバイト文字列である'\x9c'とセットで 1 つの文字列だからである。つまり'\xe5\xa4\x9c'で「夜」となる。

このように日本語は複数セットで 1 つの文字列みたいなトークン分割がされるので、これらを考慮してデコードする必要がある。

分割されたまま復元する

シンプルに「デコードして失敗したら次のバイト文字列を連結してデコードする」という処理をすればいい。

# トークン分割後のバイト文字列を、分割されたまま復元して返す
# ex.) [b'\xe3\x81\x93\xe3\x82\x93\xe3\x81\xab', b'\xe5\xa4', b'\x9c'] -> ["こんに", "夜"]
def decodeByteStrings(byte_strings) -> List[str]:
    temp_byte = b""
    result = []
    for b in byte_strings:
        try:
            temp_byte += b
            decode_str = temp_byte.decode()
            temp_byte = b""
            result.append(decode_str)
        except:
            pass
    return result

以下のような感じで使っている

# トークン数とトークン分割文字列を取得
def get_token():
    text = "こんに夜"
    enc = tiktoken.get_encoding("cl100k_base")
    tokens = enc.encode(text)
    bs = enc.decode_tokens_bytes(tokens)
    token_text = decodeByteStrings(bs)
    return {"token": len(tokens), "token_text": token_text}

# {
#   "token": 3,
#   "token_text": [
#     "こんに",
#     "夜"
#   ]
# }