Elixir（Livebook）でOpenCALMを動かす

Elixir（Livebook）でOpenCALMを動かしてみます。

依存ライブラリの準備

記事執筆現在のBumblebee（0.3.0）では、GPT-NeoXに対応したコードがリリースされていないので、GitHubリポジトリから取得するようにします。

Add GPT-NeoX (StableLM) by seanmor5 · Pull Request #204 · elixir-nx/bumblebee

Mix.install(
  [
    {:bumblebee, github: "elixir-nx/bumblebee"},
    {:exla, ">= 0.0.0"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

モデルとトークナイザの準備

上記の準備さえできれば、あとは他のモデル同様にやっていくだけです。

model = "cyberagent/open-calm-7b"

{:ok, model_info} = Bumblebee.load_model({:hf, model})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, model})

テキストを生成する

Bumblebee.Text.generationを使うだけで動きます（7bモデルをM1 MacのCPUで動かしてみて、20〜30分ぐらいかかりました）。

{:ok, generation_config} = Bumblebee.load_generation_config({:hf, model})

Bumblebee.Text.generation(model_info, tokenizer, generation_config)
|> Nx.Serving.run("大規模言語モデルは、")
|> then(&(&1.results))

[
  %{
    text: "大規模言語モデルは、コンピュータ・シミュレーションによって、非線形な現象を、非線形な関数で表現する。非線形"
  }
]

途中で切れてしまいましたが、それっぽいテキストが生成されました。

履歴付きの会話をする

Bumblebee.Text.conversationも使ってみたいところです。

# `generation_config`は上記と同じ
serving = Bumblebee.Text.conversation(model_info, tokenizer, generation_config)

history = nil
message = "こんにちは！"

%{text: text, history: history} =
  serving
  |> Nx.Serving.run(%{text: message, history: history})

message = "調子どう？"

%{text: text, history: history} =
  serving
  |> Nx.Serving.run(%{text: message, history: history})

%{
  text: " トップページ > 債務整理 自己破産 デメリット\n債務整理 自己破産 ",
  history: [
    generated: " トップページ > 債務整理 自己破産 デメリット\n債務整理 自己破産 ",
    user: "調子どう？",
    generated: " トップページ > 静岡県 > 富士市|債務整理の無料相談は法テラス",
    user: "こんにちは！"
  ]
}

よくわからないことになってしまいました。

おわりに

使ってみた、ぐらいの感じでしたが、Elixir（Livebook）でOpenCALMを動かしてみました。HuggingFaceにあるモデルを簡単に試すことができて、Bumblebeeすごいですね。