この記事は、LLM・LLM活用 Advent Calendar 2024 の 10 日目の記事になります。

https://qiita.com/advent-calendar/2024/large-language-model

はじめに

テキストを生成するモデルである大規模言語モデル (LLM) は transformers などのライブラリで簡単に扱えるようになりました。

また、近年はテキストだけでなく、画像や動画をもとにテキストを生成できるマルチモーダルなモデルも増えており、その際にテキストや画像を機械学習モデルへ入力するための前処理は段々と複雑になってきています。

テキストをトークンに分割する処理や画像をパッチに分割する処理などの前処理担当を、プリプロセッサーやプロセッサーと呼びますが、transformers では、テキストと画像を一緒に前処理するプロセッサーの一例として LlavaProcessorが実装されています。

LlavaProcessor では画像認識モデル用の CIPImageProcessor と テキストトークナイズ用の LlavaTokenizer をまとめて扱うことができますが、このようなプロセッサーをカスタムする方法はあまり紹介されていないので紹介したいと思います。

今回は、特殊な例として 二つの異なるトークナイザーを一つのプロセッサーで扱う方法 と、そのプロセッサーを簡単に push_to_hub() や from_pretrained() で 保存・読み込みできるように する方法を紹介します。この方法を応用することで、好きなだけトークナイザーやプロセッサーをまとめて管理することができるようになります。

カスタムコードで独自のモデルを transformers で扱う方法については、LLM・LLM活用 Advent Calendar 2024 の9日目の記事 Huggingface Transformersに自分のモデルを追加してみた！@weak_kajuma を先に読んでおくと良いかもしれません。

プロセッサーの例

二つのトークナイザーを内包するプロセッサーの例を示します。

実際に動作する例を HuggingFace Hub にアップしています:

https://huggingface.co/p1atdev/multi-tokenizers-processor-sample

ディレクトリ構成

./
├── models
│ ├── __init__.py
│ └── processor_multi.py <- プロセッサー本体
...

models/__init__.py は空ファイルですが、カスタムコード登録の関係で必要になります。

実装

models/processor_multi.py では MultiProcessor という名前でプロセッサーを実装しています。これは、基本的に LlavaProcessor と同様の処理になっていますが、複数のトークナイザーを扱うために幾つか変更 が加えられています。

コメントにも書いてありますが、実装されているクラス・関数はそれぞれ次のような役割です:

• MultiProcessorKwargs: トークナイザーに渡すデフォルト引数を定義
  • LlavaProcessor で使われているデフォルト引数の扱いと若干異なりますが、最終的に同じことが実現できればいいのでここの形式は重要ではないです
• MultiProcessor: プロセッサー本体のクラス
  • 変数
    • attributes:　プロセッサーが持つ子トークナイザの名前を指定
    • valid_kwargs: (よくわかってない) おそらく前処理実行時に受け取れる引数名？
    • tokenizer_1_class, tokenizer_2_class: それぞれのトークナイザーのインスタンス化に使うクラス名。
    • tokenizer_1, tokenizer_2: それぞれのトークナイザーのインスタンス
      • それぞれのトークナイザーのクラス指定には attributes で指定した名前に揃える必要があります
  • 関数
    • __init__: プロセッサーの初期化
    • __call__: processor(text_1="テキスト", text_2="テキスト") で呼び出されるときの関数。ここで受け取れる引数を指定&前処理実行
    • batch_decode: 複数のシーケンスをまとめてデコードする関数
      • ここでは二つ目のトークナイザーだけでデコードしています
    • decode: 一つのシーケンスをデコードする関数
      • こちらも同様に二つ目のトークナイザーだけでデコードしています
    • model_input_names: (よくわかってない) おそらく __call__ するときに受け取れる引数名？
    • _get_arguments_from_pretrained: save_pretrained する際に呼ばれる、子トークナイザーをディスクに保存する処理
    • save_pretrained: ローカルにプロセッサーを保存する処理。通常の実装では複数のトークナイザーを持つと正しく保存できないので、修正した処理に変更しています

使用例

以下の二つのトークナイザーを持たせてみます

• 一つ目のトークナイザー (tokenizer_1): llm-jp/llm-jp-3-1.8b のトークナイザー
• 二つ目のトークナイザー (tokenizer_2): Qwen/QwQ-32B-Preview のトークナイザー

これらのトークナイザーは、AutoTokenizer で読み込めて PreTrainedTokenizer として扱えるのならば何でも使えます。

コードではこのようになります：

from transformers import AutoTokenizer, AutoProcessor
from models.processor_multi import MultiProcessor

# push_to_hub 用に AutoProcessor に登録
MultiProcessor.register_for_auto_class("AutoProcessor")

# プロセッサーを作成
processor = MultiProcessor(
 tokenizer_1=AutoTokenizer.from_pretrained("llm-jp/llm-jp-3-1.8b"),
 tokenizer_2=AutoTokenizer.from_pretrained("Qwen/QwQ-32B-Preview"),
)

# エンコード
print(processor(
 text_1="テキスト1",
 text_2="テキスト2",
))
# {'input_ids': tensor([[ 1, 43412, 28745]]), 'attention_mask': tensor([[1, 1, 1]]), 'input_ids_2': tensor([[56833, 61803, 70534, 17]]), 'attention_mask_2': tensor([[1, 1, 1, 1]])}

# push_to_hub で huggingface hub にアップロード
processor.push_to_hub(MY_REPO_NAME, private=True)

from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained(
 MY_REPO_NAME,
 trust_remote_code=True, # カスタムコードなので必要
)

複数のトークナイザーを持ったまま保存するための小技

transformers では、非常に簡単に save_pretrained(), push_to_hub(), from_pretrained() などで保存やhubへのアップロード、読み込みができるわけですが、裏側でどんな処理が行われているかを考えたことがあるでしょうか？

プロセッサーで push_to_hub() が呼び出さると、以下の図のような処理が行われます。

attributes や attribute は、コード中におけるプロセッサーが持っている子トークナイザー・プロセッサーのことです。

push_to_hub() は大まかな流れとして、 save_pretrained() してから huggingface_hub.hf_api を用いて、保存されたファイルをすべて HuggingFace Hub にアップロードするという形になります。

そのため、適切に push_to_hub するためには save_pretrained でローカルディスクに書き込む処理が正常に行えれば良いことがわかります。

プロセッサーにおける from_pretrained は以下ののうになります:

ただし、既存の ProcessorMixin は現時点 (2024/12/04) では、 save_pretrained、from_pretrained する際に サブディレクトリを指定せず、保存ディレクトリ直下にすべてを保存してしまう ため、一つのプロセッサーが複数のトークナイザー・プロセッサーを持つと同じディレクトリに保存・上書きしてしまうため、競合してしまいます。実際に保存処理を修正せずに save_pretrained、from_pretrained をすると、二つとも同じトークナイザー(一番最後に保存された方)が使われてしまいます。

そこで、保存・読み込み処理を少し調整し、それぞれの子トークナイザーをサブディレクトリに保存するように変更することで、好きなだけたくさんのトークナイザー持ったプロセッサーを作成できるようになります。トークナイザーだけでなく画像のプロセッサーについても同様です。

注意点

カスタムコードを Huggingface Hub にアップする際に AutoProcessor に register する必要があるわけですが、この登録処理では クラス名に依存して処理が若干変わる罠 があり、それによって適切に from_pretraiend できなくなることがあります。

具体的には、 save_pretrained の途中で呼び出される custom_object_save 関数が保存するカスタムプロセッサーのクラス名に基づいた処理をしており、Tokenizer が名前に含まれていると auto_map に一般のトークナイザー用のオプションである slow_tokenizer の枠を用意してしまい、from_pretrained する際に auto_class への引数に余計なものが増えて正常に読み込めなくなります。 そのため、カスタムプロセッサーのクラス名には tokenizer を含めないようにする必要があります。

当該処理:

https://github.com/huggingface/transformers/blob/329f5dbf97a5cb2473914c88c05aa3dcb242e19a/src/transformers/dynamic_module_utils.py#L583-L593

まとめ

本記事では transformers ライブラリのカスタムプロセッサーとして、複数のトークナイザーを競合しないように扱う方法を紹介しました。

今後の LLM・LLM活用 Advent Calendar 2024 もお楽しみください。

https://qiita.com/advent-calendar/2024/large-language-model

関連情報

https://huggingface.co/docs/transformers/v4.47.0/en/custom_models