Hugging Face Optimumでrinna GPTモデルをONNXに変換する

ONNXランタイムは、モバイルやエッジを含むクロスプラットフォームの推論エンジンです。ONNXモデルに変換することで、モデルを様々な環境で動かすことができます。

huggingface/optimumを利用することで、Transformersのモデルを簡単にONNXに変換できます。

本稿はだいたいHugging Faceの公式ドキュメントに書いてあることです。

環境

• Python 3.10.12
• optimum 1.8.8

今回はGoogle Colabを利用しました。予め!pip install optimum[exporters]でインストールしておきます。

変換

rinna社の提供する日本語GPTモデルであるrinna/japanese-gpt2-xsmallをONNXに変換してみましょう。下記のコマンドで、my_onnx_gpt/というディレクトリに変換されたモデルが出力されます。

!optimum-cli export onnx --model rinna/japanese-gpt2-xsmall my_onnx_gpt/

以上です。お疲れ様でした。

念のため確認してみましょう。

!ls my_onnx_gpt/

出力

config.json		   decoder_with_past_model.onnx  spiece.model
decoder_model_merged.onnx  generation_config.json	 tokenizer_config.json
decoder_model.onnx	   special_tokens_map.json	 tokenizer.json

どうやら.onnxに変換されているらしいことが分かります。

実行

optimumを利用してONNXモデルを実行するには、普通のTransformersを利用して実行するときとあまり変わりません。Hugging Faceの公式ドキュメントから、適切なモデルクラスを選択して差し替えてください。今回であればORTModelForCausalLMです。

from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForCausalLM

MODEL_NAME = "my_onnx_gpt"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = ORTModelForCausalLM.from_pretrained(MODEL_NAME)

onnx_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)

gen = onnx_gen("昔々あるところに")
print(gen)

出力※毎回変化します

[{'generated_text': '昔々あるところに突然現れた - Cの C と言えば、C のキャプチャーを撮影したのが本当らしい。 そのカメラは「 500」という名称で、「'}]