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

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

注: 現在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](https://thegeekpage.com/make-windows-11-accept-file-paths-over-260-characters/)

そして、@suzuki_sh様の["Hugging Face Optimumでrinna GPTモデルをONNXに変換する"](https://qiita.com/suzuki_sh/items/7c0fb56f24948ea8e6ed)の通りに実行します。

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

- python test code
[link](https://github.com/mori32/codesnipets/blob/main/onnx-test/onnx-verify.py)
```
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
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3540211/73fd6f88-1ede-3c83-fc67-1869bd5c80c8.png)

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

[vcpkg.jon](https://github.com/mori32/codesnipets/blob/main/onnx-test/vcpkg.json), [vcpkg-configuration.json](https://github.com/mori32/codesnipets/blob/main/onnx-test/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](https://github.com/google/sentencepiece/blob/master/doc/api.md)を参考にします。単にトレーニング済みのモデルを使うには、`SentencePieceProcessor::Load()`, `Encode()`, `Decode()` だけ使えばよいようです。
モデルファイルはは、onnxコンバートフォルダの`my_onnx_gpt/spiece.model`です。

自作wrapperクラスは、[こちら](https://github.com/mori32/codesnipets/blob/main/onnx-test/src/tokenizer.cpp)。ただしonnx側がtoken配列にint64_tの配列を要求しているので、このクラスで`vector<int64_t>`にしています。

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

# onnxモデルを確認する

まずは、onnxモデルの入出力を確認します。[netron](https://netron.app/)という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](https://learn.microsoft.com/ja-jp/uwp/api/windows.ai.machinelearning?view=winrt-22621) です。ですが、このリファレンスだけでは使い方がわかりづらいので、どこかのサイトでほかの方の作ったサンプルコードを参照にしたはずですが、失念しました。思い出したら更新します。

手順の概要は以下になりますが、[サンプルコード](https://github.com/mori32/codesnipets/blob/main/onnx-test/src/onnxConnector.cpp)を見た方が分かりやすいかもしれません。
- 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リストの最後に追加して、再度モデルを実行して、予測テキストを伸ばしていきます。

[サンプルプログラム](https://github.com/mori32/codesnipets/tree/main/onnx-test)は、私のところでは以下を出力します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3540211/d6b2d4a1-304d-cff8-ad91-8c4aeb1e639e.png)

# まとめ

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