rinna GPTモデルをonnxに変換し、C++から使ってみる

GPTを使うサンプルコードのほとんどは、pythonで書いてあり、試してみるには十分なのですが、実際のプログラムに組み込むには多少困難になります。この記事では、rinna GPTモデルをonnxモデルに変換し、C++から使ってみます。
サンプルコードの GitHub repo はこちらのです。

注: 現在2023/8/30ですが、半年前に同じこと試したときから、いくつか変わっている点がありました。したがって今後も、このままでは動かなくなることもあるでしょう。

環境

• Windows 11
• Visual Studio 2022
• Python 3.11.5

rinnaモデルをonnxに変換

optimumのintsallに、Windows ロングファイ名対応に設定する必要がありました。以下の記事に倣いました。Make Windows 11 Accept File Paths over 260 Characters

そして、@suzuki_sh様の"Hugging Face Optimumでrinna GPTモデルをONNXに変換する"の通りに実行します。

途中でWarningがたくさん出ますが、pythonのテストコードが実行されればよしとします。

• python test code
  link

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)

• output
  

sentence pieceをvcpkgでインストール

sentence pieceは、rinna GPTモデルが使っているtokenizerライブラリです。Visual Studioからは、vcpkgでライブラリをインストールできるので、やってみます。vcpkgは、Microsoftが提供するC/C++ライブラリ向けのパッケージ管理システムです。

vcpkg.jon, vcpkg-configuration.jsonを準備
vcpkg

{
  "$schema": "https://raw.githubusercontent.com/microsoft/vcpkg-tool/main/docs/vcpkg.schema.json",
  "dependencies": [
    "sentencepiece"
  ]
}

vcpkg-configuration.json

{
  "default-registry": {
    "kind": "git",
    "baseline": "638b1588be3a265a9c7ad5b212cef72a1cad336a",
    "repository": "https://github.com/microsoft/vcpkg"
  },
  "registries": [
    {
      "kind": "artifact",
      "location": "https://github.com/microsoft/vcpkg-ce-catalog/archive/refs/heads/main.zip",
      "name": "microsoft"
    }
  ]
}

sentence pieceはデフォルト設定では対応していないようなので、下記のコマンドをvcpkg.jsonと同じフォルダで実行し、インストールします。

vcpkg install --triplet x64-windows-static

sencence pieceを使う

sentence piece API docを参考にします。単にトレーニング済みのモデルを使うには、SentencePieceProcessor::Load(), Encode(), Decode() だけ使えばよいようです。
モデルファイルはは、onnxコンバートフォルダのmy_onnx_gpt/spiece.modelです。

自作wrapperクラスは、こちら。ただしonnx側がtoken配列にint64_tの配列を要求しているので、このクラスでvector<int64_t>にしています。

これで、日本語文字列とtoken ID配列の相互変換ができました。

onnxモデルを確認する

まずは、onnxモデルの入出力を確認します。netronというonnxモデルのvisualizerを使い、my_onnx_gpt/decoder_model.onnxを開きました。すると入力は

name: input_ids
tensor: int64[batch_size,sequence_length]

name: attention_mask
tensor: int64[batch_size,sequence_length]

出力は

name: logits
tensor: float32[batch_size,sequence_length,32000]

とわかります。特にDimension数は確認しないと、パラメータを作れません。

onnxモデルをWindows.AI.MachineLearning APIで使う

onnxモデルAPIを、Windowsは提供しています。Windows.AI.MachineLearning です。ですが、このリファレンスだけでは使い方がわかりづらいので、どこかのサイトでほかの方の作ったサンプルコードを参照にしたはずですが、失念しました。思い出したら更新します。

手順の概要は以下になりますが、サンプルコードを見た方が分かりやすいかもしれません。

• LearningModel::LoadFromFilePath() でモデルインスタンスを作ります
• LearningModelSession を LearningModel から作ります
• LearningModelBinding を　LearningModelSession　から作ります。
• LearningModelBinding に "input_ids" の配列を作りバインドします。
• LearningModelBinding に "attention_mask" の配列を作りバインドします。
• LearningModelSession::Evaluate() で計算し、LearningModelEvaluationResult を取得します。
• LearningModelEvaluationResult::Outputs プロパティから "logits" を取り出します。
• (とりあえず Greedy を実装するので) logits の中から最大の token ID を見つけて、次 token とします。

ちなみ logits の3次元配列ですが、batch は入力が一つなので出力も一つです。
sequence_length は、入力 token の sequence の長さと同じです。32000 の配列が、次の token の logits になります。(たぶん)

	token 1	token 2	toekn 3
input text	昔々	ある	ところに
logits	token 2	token 3	token 4

と、次のtokenのlogitsが入っているので、最後のtokenのlogitsを評価します。

あとは、最大のtoken IDを、今までのtokenリストの最後に追加して、再度モデルを実行して、予測テキストを伸ばしていきます。

サンプルプログラムは、私のところでは以下を出力します。

まとめ

rinna GPTモデルをonnxフォーマットに変換して、C++から使ってみました。
ですが、Greedy アルゴリズムがこれでいいか、ORTModelForCausalLM と同じにするにはどうすればいいか、まあ謎はいっぱいです。