Published on

tokenizerの学習

Author
株式会社イエローバックの機械学習エンジニアです

はじめに

transformers v4.9 で導入された以下の機能を試してみます。

既存のトークナイザの設定をベースに新たにトークナイザを学習できるようになりました。

今回は gpt2 をベースにして日本語特化のトークナイザを作成してみます。

必要な python packages

  • datasets: 1.10.2
  • transformers: 4.9.0

学習

tokenizer.train_new_from_iterator()にデータを渡してあげると tokenizer と同じ設定で学習してくれます。

今回使用する corpus は mc4 の日本語データとします。 かなり大きなデータセットなので先頭の 10 万ドキュメントのみ利用します。 mc4 は huggingface の datasets を使って容易に取得できます。datasets 1.10 で導入された streaming モードを使うと、 全データのダウンロードを待つ必要がなくなりますので、この機能を利用します。

train_tokenizer.py
from datasets import load_dataset
from transformers import AutoTokenizer
from more_itertools import chunked

tokenizer = AutoTokenizer.from_pretrained("gpt2")

dataset = load_dataset("mc4", "ja", streaming=True, split='train')
ds_sub = dataset.take(100000)
corpus = chunked((x['text'] for x in ds_sub), 1000)
new_tokenizer = tokenizer.train_new_from_iterator(corpus, vocab_size=32000)
new_tokenizer.save_pretrained('new_tokenizer')

Core i9 9900K のマシンを用いて 20 分程度で終了しました(コーパスのダウンロード時間は除く)。 ディレクトリ new_tokenizerには新たに学習された tokenizer ができています。

$ ls -l new_tokenizer/
total 2468
-rw-rw-r-- 1 ono ono  505721 Jul 23 20:07 merges.txt
-rw-rw-r-- 1 ono ono      90 Jul 23 20:07 special_tokens_map.json
-rw-rw-r-- 1 ono ono 1289503 Jul 23 20:07 tokenizer.json
-rw-rw-r-- 1 ono ono     236 Jul 23 20:07 tokenizer_config.json
-rw-rw-r-- 1 ono ono  719587 Jul 23 20:07 vocab.json

サンプルテキストを用いて id 列に変換してみます。

train_tokenizer.py
sample_text = '今回はgpt2をベースにして日本語特化のトークナイザを作成してみます。'

print(tokenizer.encode(sample_text))
[20015, 232, 32368, 252, 31676, 70, 457, 17, 31758, 35604, 6312, 8943, 28618, 22180, 28134, 33768, 98, 17312, 105, 45739, 252, 31965, 117, 44293, 244, 5641, 13298, 42869, 26229, 11482, 48458, 31758, 43291, 22755, 238, 22180, 28134, 2515, 123, 30159, 33623, 16764]

print(new_tokenizer.encode(sample_text))
[3433, 71, 4729, 18, 21251, 2103, 5214, 875, 25586, 10570, 4223, 801, 10349, 2534, 298, 257]

このサンプルでは、オリジナル版は 42 トークンに対して日本語特化版では 16 トークンとなりました。

まとめ

tokenizer.train_new_from_iterator() を用いて日本語特化の gpt2 トークナイザが容易に作成できました。