chiTra モデルはワークス徳島人工知能 NLP 研究所が提供する事前学習済み BERT モデルで、 Hugging Face の Transformers ライブラリを使用して作成されています。

Transformers には名称やモデルファイルからインスタンスを自動で読み込む Auto Classes という機能があり、評価実験で複数のモデルを使い分ける際など大変便利です。
しかし、chiTra トークナイザは独自実装である関係上、そのままでは読み込むことができません。

本記事ではこの機能で chiTra トークナイザを読み込めるようにすることをゴールとして、手法を二つ解説します。

情報元

本記事は主に公式ドキュメントの Sharing custom models で述べられている手法を chiTra を対象に書き換えたものです。

また Megagon Labs の寺田氏も類似内容の記事を書いていらっしゃいます。

0. Auto を使わずに読み込む

Auto 機能を使う話の前に、使わない場合の読み込み方について述べておきます。

この場合はクラスを明示すればよく、コードは以下のようになります。

from sudachitra import BertSudachipyTokenizer
tokenizer = BertSudachipyTokenizer.from_pretrained("path/to/chiTra-1.0")

また chiTra モデル（およびコンフィグ）については BERT アーキテクチャを採用しているため、
Auto 機能のパスからのクラス推定を利用できます。
具体的にはモデルデータを bert/ ディレクトリの内部に配置することで、以下のように Auto 機能で読み込むことができます。

from transformers import AutoConfig, AutoModel
config = AutoConfig.from_pretrained("path/to/bert/chiTra-1.0")
model = AutoModel.from_pretrained("path/to/bert/chiTra-1.0")

1. 読み込み用コードを追加する

一つ目の方法はモデルデータを改変し、Auto 機能で使用されるコードを追加するものです。

トークナイザが読み込まれる際にはモデルデータ内の tokenizer_config.json が参照されます。
ここに以下のように auto_map 項目を追加します。

tokenizer_config.json

{
    # ... 前略
    "tokenizer_class": "BertSudachipyTokenizer",
    "auto_map": {
      "AutoTokenizer": [
          "modeling.BertSudachipyTokenizer",
          null
      ]
   }
}

null について

上記の変更では AutoTokenizer に modeling.BertSudachipyTokenizer と null の二つを与えています。
Hugging Face Tokenizers のトークナイザには Slow と Fast の二つの区分があり、ここではそれぞれに対応するクラスを登録する形式になっているためです。

chiTra トークナイザは Fast 形式に対応していないため、二つ目の枠については null を与えています。

これによって AutoTokenizer を使用した際に、モデルデータにある modeling モジュールが読み込まれ、
BertSudachipyTokenizer の探索範囲に追加されるようになります。
従ってさらに以下内容の python コードをモデルディレクトリに追加することで、トークナイザを Auto 機能で読み込めるようになります。

modeling.py

from sudachitra import BertSudachipyTokenizer

この追加コードでは任意の処理を実行できてしまうため、読み込みの際に引数 trust_remote_code=True を与えて実行を明示的に許可する必要があります。またセキュリティの観点から、外部のコードをこの方法で参照する場合は内容に目を通した上で revision まで指定することが強く推奨されています（参照）。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("path/to/chiTra-1.0", trust_remote_code=True)

2. Auto 機能へクラスを登録する

Auto 機能を実装する各クラスには、カスタムクラスを登録するための register メソッドが用意されています。
これを利用するのが二つ目の方法です。

この方法では、AutoTokenizer を使う python ファイル内で事前に登録を行います。

from sudachitra import BertSudachipyTokenizer
from transformers import AutoTokenizer

# クラスを登録
AutoTokenizer.register(BertSudachipyTokenizer.__name__, slow_tokenizer_class=BertSudachipyTokenizer)

tokenizer = AutoTokenizer.from_pretrained("path/to/chiTra-1.0")

この方法では Auto 機能を使う箇所ごとに登録のコードを追加する必要がありますが、
代わりに AutoTokenizer の引数変更は必要ないという利点があります。

まとめ

以上、AutoTokenizer から chiTra トークナイザを読み込む方法二つを紹介しました。

個人的には chiTra を扱う際には

データを bert/ 以下に配置してモデルを読めるようにしつつ、
二つ目の AutoTokenizer.register をコードに追加する方法を使う

とするのがモデルデータの変更が不要な分楽なのかな、と考えています。

ご自分の環境や手法にあった方法で chiTra を使ってみてください。

AutoTokenizer で chiTra トークナイザを読み込む

情報元

0. Auto を使わずに読み込む

1. 読み込み用コードを追加する

2. Auto 機能へクラスを登録する

まとめ