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

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

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

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

# 情報元

本記事は主に公式ドキュメントの [Sharing custom models](https://huggingface.co/docs/transformers/main/en/custom_models) で述べられている手法を chiTra を対象に書き換えたものです。

また Megagon Labs の寺田氏も[類似内容の記事](https://www.megagon.ai/jp/blog/autotokenizer/)を書いていらっしゃいます。

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

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

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

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

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

```python
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` 項目を追加します。

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

<details><summary>null について</summary><div>

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

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

</div></details>

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

```python:modeling.py
from sudachitra import BertSudachipyTokenizer
```

この追加コードでは任意の処理を実行できてしまうため、読み込みの際に引数 `trust_remote_code=True` を与えて実行を明示的に許可する必要があります。またセキュリティの観点から、外部のコードをこの方法で参照する場合は内容に目を通した上で revision まで指定することが強く推奨されています（[参照](https://huggingface.co/docs/transformers/main/en/custom_models#using-a-model-with-custom-code)）。

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

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

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

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

```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 を使ってみてください。