Rubyで文章生成AI「GPT-2」を実行して遊んでみた #Ruby

はじめに

Python 以外の言語から AI にアクセスするのは難しすぎると考えている人が多いと思います。しかし、最近は、事前トレーニング済みのモデルが多数公開されており、さまざまなプラットフォームやプログラミング言語から AI 推論を簡単に実行できるようになっています。ここでは、Ruby 言語から ONNX ランタイムを実行して英文を生成できる GPT-2 モデルで遊んでみたいと思います。 ONNX ランタイムを使用すると、推論は簡単に実行できます。

事前準備とインストール。

下記の３つのGemをインストールします。

onnxruntime - ONNX Runtime のRubyバインディングです。
tokenizers - Hugging Face の提供しているTokenizerのRubyバインディングです。
numo-narray - Rubyの行列計算ライブラリです。NumPyと同じようなことができます。

gem install onnxruntime
gem install tokenizers
gem install numo-narray

onnxruntime gem は内部にプラットフォームごとのバイナリファイルを梱包済みですので、ONNX Runtimeを別途インストールする必要はなく、gemをインストールするだけで動作します。また、単数形の tokenizer は全く別のGemですのでご注意ください。

なお、tokenizersの拡張機能はRustで書かれているためインストールにはRustが実行できる環境が必要です。Windowsで動かしたい場合やRustがない環境では、tokenizersのかわりにblingfireを使います。実装はこんなかんじです。

ONNX モデルの入手

ここではONNX公式が配布するGPT-2のモデルを利用します。リンク先からGPT-2-LM-HEADをダウンロードします。

上の画面で右クリックでダウンロードしようとするとリンク先はまだhtmlなので失敗します。一度リンクを踏んで、次のページの右上のDownloadボタンからモデルをダウンロードしてください。

gpt2-lm-head-10.onnx をスクリプトと同じディレクトリに配置して下さい。Netronというソフトウェアを使うと、ONNXモデルの中身を閲覧することができます。

コードを書く

準備は完了しました。それでは実際に動かすRubyのコードを準備しましょう。

require "tokenizers"
require "onnxruntime"
require "numo/narray"

s = "Why do cats want to ride on the keyboard?"

# トークナイザーを準備する。ここではgpt2を指定
tokenizer = Tokenizers.from_pretrained("gpt2")

# モデルを読み込む
model = OnnxRuntime::Model.new("gpt2-lm-head-10.onnx")

# 文字列を行列に変換
ids = tokenizer.encode(s).ids

# 単語を10回取得
10.times do
  o = model.predict({ input1: [[ids]] })
  o = Numo::DFloat.cast(o["output1"][0])
  ids << o[true, -1, true].argmax
end

# 結果を出力：行列から文字列へ
puts tokenizer.decode(ids)

下記の部分は、Rubyとしては見慣れない記法だと思います。これは numo-narray という NumPy と同様のことができるライブラリを利用しています。詳しい使い方を知りたい方は、NArrayの簡単なつかい方をご覧ください。

  o = Numo::DFloat.cast(o["output1"][0])
  ids << o[true, -1, true].argmax

※ 実は行列操作の部分は、私も何をしているのか本質的にはわかっていなくて、ひたすらNumPyのコードをnumo-narrayに直訳しています。けれども、逆に言うと、あまりAIや数学の知識がなくても、公開されているAIのモデルは利用はできてしまいます。自動車の整備ができなくても、自動車が運転できるのと同じです。

実行してみる

ruby コマンドでも実行できますし、irb に上のコードをコピペしても動きます。

Why do cats want to ride on the keyboard?

The answer is that they do.

なんでネコはキーボードに乗りたがるの？
その答えは「そうしたいから」。

と表示されたら成功うまくいっています。

いかがでしたか？

今回実行したのは、Web上のAPIを呼び出しているわけではなく、実際に手元のマシンでDeep Learningを実行した内容になります。思ったより簡単にRubyからDeep Learning のモデルを呼び出せたかと思います。

それでは楽しいRubyライフを。この記事は以上です。