AutoTokenizerでCustom Tokenizerを読み込む方法

AutoTokenizerでCustom Tokenizerを読み込む方法

Megagon Labsで自然言語処理技術の開発パートナーとしてGiNZAの機能実装を担当している寺田です。

Hugging Faceのtransformersライブラリでは、使用するmodelやtokenizerのインスタンスをモデルファイルから自動で生成するAuto Classesという機能があります。 異なるモデルを統一的に扱う上でとても便利な機能なのですが、transformersライブラリ外のCustom Tokenizerを読み込むためにはひと工夫必要だったのでこちらでご紹介します。 例として、GiNZAでのCustom Tokenizerの扱い方を示していますが、同様の問題は他の様々な応用でも生じていると思います。

課題認識 – GiNZAを題材に

GiNZAではモデルの訓練から推論までspaCyを利用しています。
spaCyではconfigファイルを用意することでノーコードでモデルの訓練を行うことができますが、ノーコード故に融通の効かない部分もあり、その一つとしてAuto Classesでは読み込めないmodelやtokenizerが使えないという問題があります。

ginza-transformers と課題点

spaCyでhuggingface/transformersを利用するためのライブラリとしてspacy-transformersがあります。
spacy-transformers内部では以下のhuggingface_from_pretrained関数内で以下のようにAuto Classesを用いてmodel をロードしています。

def huggingface_from_pretrained(
    source: Union[Path, str], tok_config: Dict, trf_config: Dict
) -> HFObjects:
    """Create a Huggingface transformer model from pretrained weights. Will
    download the model if it is not already downloaded.
    source (Union[str, Path]): The name of the model or a path to it, such as
        'bert-base-cased'.
    tok_config (dict): Settings to pass to the tokenizer.
    trf_config (dict): Settings to pass to the transformer.
    """
    if hasattr(source, "absolute"):
        str_path = str(source.absolute())
    else:
        str_path = source
    tokenizer = AutoTokenizer.from_pretrained(str_path, **tok_config)
    vocab_file_contents = None
    if hasattr(tokenizer, "vocab_file"):
        with open(tokenizer.vocab_file, "rb") as fileh:
            vocab_file_contents = fileh.read()
    trf_config["return_dict"] = True
    config = AutoConfig.from_pretrained(str_path, **trf_config)
    transformer = AutoModel.from_pretrained(str_path, config=config)
    ops = get_current_ops()
    if isinstance(ops, CupyOps):
        transformer.cuda()
    return HFObjects(tokenizer, transformer, vocab_file_contents)

GiNZAのtransformers対応モデルであるja-ginza-electraではtokenizerとしてSudachiTraElectraSudachipyTokenizerを用いていますが、sudachitraはtransformers標準のAuto Classesには登録されていないため、ただ単にElectraSudachipyTokenizertokenizer_config.jsonで指定するだけでは読み込んではくれません。spaCyではコードは基本的にコマンドラインからpython -m spacyで起動する形になるため、import sudachitraを実行することは困難です。

これを解決するために、GiNZAではginza-transformersというライブラリを作成し、上記のhugging_face_from_pretrained関数の機能を代替していますが、spacy-transformersに対して密結合な実装にならざるを得ず、メンテナンス性の低さが課題となっています。

AutoTokenizer で ElectraSudachipyTokenizer を読み込む

AutoTokenizer.register

Auto Classesのドキュメントではregister という新規クラスをAuto Classesに登録する関数が紹介されています。

from transformers import AutoConfig, AutoModel

AutoConfig.register("new-model", NewModelConfig)
AutoModel.register(NewModelConfig, NewModel)

一見これでうまくいきそうですが、これは新しい名前を新しいクラスに割り振るためのもので、既に登録されている名前、クラスに対しては利用することができません。 今回のケースではElectraConfigに対してはライブラリ内でElectraTokenizerが既に割り当てられているためElectraSudachipyTokenizerで上書きするといったことはできない仕様になっています。

auto_map

うまくいくやり方は、tranformersでCustom Classを扱う方法としてもう一つ提供されているauto_mapを利用する方法です。

"auto_map": {
    "AutoTokenizer": [
        "module.slow_tokenizer_class",
        "module.fast_tokenizer_class"
    ]
}

上記のようなauto_maptokenizer_config.jsonに追加することで、moduleを自動でimportして利用する仕組みが利用できるようになります。

if use_fast and tokenizer_auto_map[1] is not None:
    class_ref = tokenizer_auto_map[1]
else:
    class_ref = tokenizer_auto_map[0]

module_file, class_name = class_ref.split(".")
tokenizer_class = get_class_from_dynamic_module(
    pretrained_model_name_or_path, module_file + ".py", class_name, **kwargs
)

上記設定の上でAutoTokenizer.from_pretrainedを呼び出す際には、trust_remote_codeを指定する必要があります。 またこの方法ではmodeling.py内の任意のコードが実行されてしまうため、自社開発以外のレポジトリを参照する際には、レポジトリの現状のコードにきちんと目を通した上で、revisionまで指定することが推奨されています。

revision (`str`, *optional*, defaults to `"main"`):
    The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
    git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
    identifier allowed by git.
trust_remote_code (`bool`, *optional*, defaults to `False`):
    Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
    should only be set to `True` for repositories you trust and in which you have read the code, as it will
    execute code present on the Hub on your local machine.

こちらの方法ではAutoConfigconfig.jsonを読み込めることが前提となっています。AutoConfigconfig.jsonを読み込めない場合にはregisterを用いる方法でConfigを新たな名前で登録するとよいでしょう。

結論

以下の三点の変更を行うことでCustom TokenizerをAutoTokenizerから呼び出せるようになりました。

  1. huggingface hubで次のようにCustom Tokenizerをimportするpythonファイルを配備する https://huggingface.co/megagonlabs/electra-base-japanese-discriminator/blob/main/modeling.py
# modeling.py
from sudachitra import ElectraSudachipyTokenizer
  1. tokenizer_config.jsonに以下のようにauto_mapフィールドを追加する
{
 "tokenizer_class": "ElectraSudachipyTokenizer",
 "do_lower_case": false,
 "do_word_tokenize": true,
 "do_subword_tokenize": true,
 "word_tokenizer_type": "sudachipy",
 "word_form_type": "dictionary_and_surface",
 "subword_tokenizer_type": "wordpiece",
 "model_max_length": 512,
 "sudachipy_kwargs": {"split_mode":"A","dict_type":"core"},
 "auto_map": {
    "AutoTokenizer": [
        "modeling.ElectraSudachipyTokenizer",
        null
    ]
 }
}
  1. from_pretrainedtrust_remote_code=Trueつきで呼び出す
AutoTokenizer.from_pretrained("megagonlabs/electra-base-japanese-discriminator", trust_remote_code=True)

おわりに

いかがだったでしょうか。 今回は既存のモデルに対して、Custom TokenizerをAutoTokenizerで読み込むための設定を解説しました。 より本質的にはhuggingface/transformersのBertJapaneseTokenizerが日本語NLPで使用されるSudachiPyやJuman++等に柔軟に対応するための拡張性の確保が必要だと思いますが、当面の回避方法として本解説がお役に立てば幸いです。

(筆者: 寺田 凜太郎 + 松田寛 / Megagon Labs Tokyo)

Tag: GiNZA, NLP, spaCy, transformers

Share:

Share on facebook
Share on twitter
Share on linkedin
Share on reddit
Share on email

More Blog Posts: