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としてSudachiTraのElectraSudachipyTokenizer
を用いていますが、sudachitraはtransformers標準のAuto Classesには登録されていないため、ただ単にElectraSudachipyTokenizer
をtokenizer_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_map
をtokenizer_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.
こちらの方法ではAutoConfig
がconfig.json
を読み込めることが前提となっています。AutoConfig
がconfig.json
を読み込めない場合にはregister
を用いる方法でConfig
を新たな名前で登録するとよいでしょう。
結論
以下の三点の変更を行うことでCustom TokenizerをAutoTokenizer
から呼び出せるようになりました。
- huggingface hubで次のようにCustom Tokenizerをimportするpythonファイルを配備する
https://huggingface.co/megagonlabs/electra-base-japanese-discriminator/blob/main/modeling.py
# modeling.py
from sudachitra import ElectraSudachipyTokenizer
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
]
}
}
from_pretrained
をtrust_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