huggingfaceのトークナイザーを学習する

はじめに

こんにちは。人材領域でレコメンドシステムの機能開発をしている羽鳥です。今日はみなさんが大好きな huggingface のライブラリを使って、自然言語処理を行う際に欠かせないトークナイザーを学習させる方法について書いていこうと思います。

huggingface

近年の自然言語処理においてはBERT¹をベースにしたモデルは欠かせないものになってきました。そしてその実装として市民権を得ているものの一つがhuggingfaceのライブラリでしょう。

huggingfaceのライブラリを使うと、いい感じに隠蔽されたインターフェイスを利用して多くのモデルを内部構造を気にすること無く使えるようになります。 Kaggleなどの分析プラットフォーム上で行われる自然言語を用いるコンペでもデファクトスタンダードとなっており、多くのユーザーに利用されています。

トークナイザー

BERTに限らず自然言語を扱うモデルに文章を食わせる際には、多くの場合文章を何らかのルールに従って分割する必要あります。モデルが見る文章の最小単位をトークンといい、文章をトークンに分けることをトークナイズと言います。

トークナイズの方法によってモデルが見る文章は大きく変わり得るため、モデリング以前の前処理工程としてはかなり重要な部分を占めています。トークナイズの方法にはいくつか種類があり、それぞれメリットとデメリットがあります。以下の3つの分割方法の説明は公式ドキュメントより内容を引用しています。

1.単語単位の分割

まずはぱっと思い浮かぶところとして単語単位で分割するという方法があります²。これはご想像の通り文に登場する1単語を1トークンとするものです。

こちらは大変簡単に実装できるというメリットがあるものの、「tokenization」と「tokenize」など意味的に似ているもの同士が全く別のトークンとして扱われてしまうという問題があります。

2.文字単位の分割

上記の問題を解決するための案として文字単位での分割をする方法があります³。これは文字通り、1文字を1トークンとするものです。

この方法では意味的に似ているものを捉えられる一方で、「apple」の文頭の「a」も「at」の文頭の「a」も同じトークンとして学習するために、全く異なる文脈で用いられるトークンを同じものとして学習してしまうという問題があります。実際に自然言語がテーマの分析コンペに参加していても、文字ベースのトークナイザーを使って学習させたモデルが既存モデルに大きく精度で勝ることは少ない気がします。

3.部分文字列単位の分割

ここまで見てきたように単語単位で分割しても、文字単位で分割してもなかなかうまくはいかないことがわかりました。それを解決するために採用されているのが部分文字列単位の分割(subword tokenization)です⁴。これは単語を部分文字列に分割し、学習対象の文章に多く現れるものをトークンとする手法です。

例えば先程の例だと「tokenization」と「tokenize」に対するトークンとして「token」が抽出されることがあります。また、この手法は単語を意味のあるトークン単位に分割できるだけでなく、トークナイザーが持つ語彙の数を小さくできるという利点もあります。このような部分文字列をユーザーが独自に設定したデータセットを用いて獲得することをトークナイザーの学習といいます。

huggingfaceのトークナイザーはユーザーが独自に定義したデータセットを用いて学習を行う機能が提供されており、以下ではその方法を見ていきたいと思います。

実装例

トークナイザーの学習は以下の手順で行うことができます。

モデルに食わせる1文章を定義する
文章を取得するイテレータを用意する
tokenizer.train_new_from_iterator()でトークナイザーの学習を行う

今回は “microsoft/graphcodebert-base” のトークナイザーを、最近Kaggleで行われた Google AI4Code – Understand Code in Python Notebooks のデータセットを使って学習する場合を考えてみます。実装例として、以下のようなコードが考えられます。

from transformers import AutoTokenizer
from typing import List

def get_training_corpus(texts: List[str]):
    for start_idx in range(0, len(texts), 1000):
        samples = texts[start_idx : start_idx + 1000]
        yield samples

tokenizer = AutoTokenizer.from_pretrained("microsoft/graphcodebert-base")

with open("./trainning_text.txt") as f:
    texts = f.readlines()

training_corpus = get_training_corpus(texts)
new_tokenizer = tokenizer.train_new_from_iterator(
    training_corpus, vocab_size=tokenizer.vocab_size
)
new_tokenizer.save_pretrained("./outputs/")

training_text.txtには事前に定義した文章が改行区切りで入っており、get_training_corpusでそれを1000件ずつのまとまりで取得しています。何をもって一つの文とするかは色々と定義が考えられるところですが、今回はコンペの設定に合わせて1つのnotebookを1つの文章として扱っています。

そのようにしてできたコーパスをtrain_new_from_iteratorに食わせて学習しています。ここでは後述の通り学習前後でvocab_sizeが変わらないように引数設定しています。

トークナイザーの学習がうまくできているか、ちょっとした例で確認してみます。

example ="""
def get_training_corpus(texts):
    for start_idx in range(0, len(texts), 1000):
        samples = texts[start_idx : start_idx + 1000]
        yield samples
"""
print(tokenizer.tokenize(example)[:10])
# ['Ċ', 'def', 'Ġget', '_', 'training', '_', 'cor', 'p', 'us', '(']

print(new_tokenizer.tokenize(example)[:10])
# ['Ċ', 'def', 'Ġget', '_', 'training', '_', 'corpus', '(', 'texts', '):']

print(len(tokenizer.tokenize(example)))
# 69

print(len(new_tokenizer.tokenize(example)))
# 48

例文として、今回使った関数定義を使ってみました。(本論とは直接関係ありませんが、“microsoft/graphcodebert-base"はプログラミング言語を対象として事前学習されたRoBERTaベースのモデルです。)

学習前のトークナイザーは以下のように「corpus」という意味ありげな単語を分割してしまっています。「Ċ」は改行を、「Ġ」は空白を表す特殊なトークンです。

['Ċ', 'def', 'Ġget', '_', 'training', '_', 'cor', 'p', 'us', '(']

一方で学習済みのトークナイザーはうまく処理していくれているようです。

['Ċ', 'def', 'Ġget', '_', 'training', '_', 'corpus', '(', 'texts', '):']

また、トークンサイズも学習前後で69から48に減っていることがわかります。これはこの一文に限った話ではなく、多くの場合でトークナイザーを学習すると分割後のトークンサイズは減少するようです。

多くのBERTモデルは入力できるトークン長が512程度に上限設定されていることが多いので、限られたトークン長のもとで入力できる情報を増やすという意味でもトークナイザーの学習は有用かもしれません。

実装時にハマりがちなポイント

学習速度について

huggingfaceのトークナイザーはRustで実装されており、かなり高速です。ただし、ときたまやたらと時間がかかるようなこともあり、issueでも報告されています⁵。私の環境でもこの現象が発生しました。実行環境に依存するバグらしく、私の場合はpythonファイルで書いていたものをJupyter Notebookで実行したところ高速に学習が進むようになりました。

vocab sizeについて

トークナイザーを学習する際にはトークナイザーが持つ語彙の大きさ（vocab size）を設定することができます。例えばベースとするトークナイザーよりも語彙の数を増やしたいような場合はこのパラメーターで調整します。

しかし、RoBERTa系のモデルではもとのトークナイザーと異なるvocab sizeを設定してしまうと後続のタスクでエラーが発生してしまうことがあります。特別な事情が無い限りはもとのトークナイザーと同じvocab sizeを設定しておくのが良さそうです。

最後に

ここまで読んで頂いてありがとうございます。近年の自然言語を扱う分析コンペではMasked Language Modelingによる事前学習や巨大なモデルを使ったファインチューニングなど事実上必須となっているテクニックがいくつかありますが、トークナイザーの学習もその一つになりえるかもしれません。この記事によってモデルの精度が1ポイントでも向上すれば幸いです。