huggingface/transformersのAutoTokenizerから学習済みSentencePiece Tokenizerを呼び出す

#はじめに
huggingface/transformersの日本語BERTモデルには、BertJapaneseTokenizerが用意されています。これはMeCabでpre tokenizeし、wordpieceかcharacter単位にtokenizeします。
しかし、BertJapaneseTokenizerはSentencePieceに対応していません。SentencePieceを使いたい場合はどうすれば良いのでしょうか。
本記事では2種類の方法を説明します。
1つ目は、SentencePieceに対応しているAlbertTokenizerを使用する方法です。非常に簡単なのでこちらの方法がおすすめです。
2つ目は、PreTrainedTokenizerFastから呼び出せるように学習済みのSentencePieceモデルを変換する方法です。手間が掛かるのでおすすめしません。

筆者が未熟なため、記事内容の誤りや記事内容より優れた方法があるかと思います。そのような箇所を見つけた際にはコメントでご指摘いただけますと幸いです。

#参考記事
1.huggingface / transformersを使って日本語BERTの事前学習を実施してオリジナルな言語モデルを作ってみる
2.SentencePieceでの日本語分かち書きをTransformersのパイプラインに組み込む

#方法1：AlbertTokenizerを使用する
非常にシンプルです。config.jsonでtokenizer_classをAlbertTokenizerに指定し、tokenizer_config.jsonを下記のようにします。
※ config.jsonでpad_token_idが正しく設定されているか確認してください
※ special tokenがデフォルト値で問題ないか確認し、適宜変更してください

tokenizer_config.json

{
    "do_lower_case": false,
    "vocab_file":"/model_dir/spiece.model",
    "unk_token":"<unk>",
    "bos_token":"<s>",
    "eos_token":"</s>",
    "pad_token":"[PAD]"
}

これを`AutoTokenizer.from_pretraind'で読み込みます。

code1

from transformers import AutoTokenizer

# 適当な例文
text = '吾輩は猫である。名前はまだない'

tokenizer = AutoTokenizer.from_pretrained('bert_dir')
output = tokenizer(text)
print(tokenizer.convert_ids_to_tokens(output['input_ids']))

result1

['[CLS]', '▁', '吾', '輩', 'は', '猫', 'である', '。', '名前', 'はまだ', 'ない', '[SEP]']

tokenizeできています。

#方法2：学習済みのSentencePieceモデルを変換する
##PreTrainedTokenizerFastに学習済みのTokenizerを読み込む
PreTrainedTokenizerFastの__init__引数にtokenizer_fileがあります。これに、huggingface/tokenizersで作成されたTokenizerファイル(json形式)を渡すと学習済みのTokenizerを読み込めます。
つまり、SentencePieceモデルをTokenizerファイルに変換できれば良さそうです。

##学習済みSentenPieceモデルの変換
huggingface/tokenizersにはsaveメソッドが実装されているため、一度huggingface/tokenizersに学習済みSentenPieceモデルを読み込みsaveメソッドを実行すれば変換できます。
huggingface/tokenizersに学習済みSentenPieceモデルを読み込むのは、参考記事2の通りに実装すれば良いので簡単です。
※ 私の環境ではsentencepiece_model_pb2.pyのダウンロードを要求されました
※ 追記：変換用の公式スクリプトがありました。ただし、BERT用に変換するものはないため、公式スクリプトを参考にして書き直したほうが良さそうです。

convert_sentencepiece_to_tokenizer.py

import argparse

from tokenizers.implementations import SentencePieceUnigramTokenizer
from tokenizers.processors import BertProcessing

def convert_sentencepiece_to_tokenizer(model_path, save_dir):
    tokenizer = SentencePieceUnigramTokenizer.from_spm(model_path)

    # BERTで使用するため、BERT用の後処理を加えます
    tokenizer.post_processor = BertProcessing(
        cls=("[CLS]", tokenizer.token_to_id('[CLS]')),
        sep=("[SEP]", tokenizer.token_to_id('[SEP]'))
    )

    tokenizer.save(save_dir+'tokenizer.json')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_path', help='model path', type=str)
    parser.add_argument('--save_dir', help='save directory path', type=str)
    args = parser.parse_args()
    convert_sentencepiece_to_tokenizer(args.model_path, args.save_dir)

これで作成したtokenizer.jsonを、PreTrainedTokenizerFastのtokenizer_fileに引き渡せば、学習済みのSentencePieceモデルを呼び出せます。
実際にできるか試してみます。

code2

from transformers import PreTrainedTokenizerFast

# 適当な例文
text = '吾輩は猫である。名前はまだない'

tokenizer = PreTrainedTokenizerFast(
    tokenizer_file='/save_dir/tokenizer.json',
    unk_token='<unk>',
    bos_token='<s>',
    eos_token='</s>',
    cls_token='[CLS]',
    sep_token = '[SEP]',
    pad_token='[PAD]',
    mask_token='[MASK]'
)

output = tokenizer(text)
print(tokenizer.convert_ids_to_tokens(output['input_ids']))

result2

['[CLS]', '▁', '吾', '輩', 'は', '猫', 'である', '。', '名前', 'はまだ', 'ない', '[SEP]']

SentencePieceでtokenizeできています。
後は、config.jsonでtokenizer_classをPreTrainedTokenizerFastに指定し、tokenizer_config.jsonを下記のようにします。
※ config.jsonでpad_token_idが正しく設定されているか確認してください

tokenizer_config.json

{
    "do_lower_case": false,
    "tokenizer_file":"/save_dir/tokenizer.json",
    "unk_token":"<unk>",
    "bos_token":"<s>",
    "eos_token":"</s>",
    "cls_token":"[CLS]",
    "sep_token":"[SEP]",
    "pad_token":"[PAD]",
    "mask_token":"[MASK]"
}

これでcode1を実行してみます。

result3

['[CLS]', '▁', '吾', '輩', 'は', '猫', 'である', '。', '名前', 'はまだ', 'ない', '[SEP]']

こちらでもtokenizeできています。

#おわりに
本記事では2種類の方法を説明しましたが、方法1でやるべきだと思います。
方法2は、SentencePieceのalgorithmがUnigramでないと実行できません。各algorithmに合わせて変換コードを実装する必要があるため、更に手間がかかります。(XLNetやALBERT用であれば、公式のスクリプトを使えば問題ありません)対して方法1はalgorithmの影響もないため、非常にシンプルです。
方法2の利点？は、うまく拡張できればJuman++などにも対応できるかもしれないことです。そこまでちゃんと調べていないため、機会があれば試してみようと思います。(できなさそうです)

(余談ですが、2種類の方法があるのは自分で方法2を考えた後に方法1を知ったためです。ドキュメントの読み込みとBERT関連の知識不足を思い知りました。)