Livebook で Bumblebee から RoBERTa を使って質問に対する回答を生成する

## はじめに

Bumblebee の公式サンプルに従ってモデルを実行してみるシリーズです

今回は RoBERTa という自然言語処理 AI モデルを利用して、 AI に質疑応答してもらいます

このシリーズの記事

- [画像分類: ResNet50](https://qiita.com/RyoWakabayashi/items/b9b4799683af5aaa29dd)
- [画像生成: Stable Diffusion](https://qiita.com/RyoWakabayashi/items/8070f6bc372015be6a8d)
- [文章の穴埋め: BERT](https://qiita.com/RyoWakabayashi/items/dcc03a8d19d7320dda8f)
- [文章の判別: BERTweet](https://qiita.com/RyoWakabayashi/items/1f8624d717fb2b1dbdc0)
- [文章の生成: GPT2](https://qiita.com/RyoWakabayashi/items/6a780c2858ec952fa1e1)
- 質疑応答: RoBERTa　(ここ)
- [固有名詞の抽出: bert-base-NER](https://qiita.com/RyoWakabayashi/items/82ac1dddfdddc8a2f79d)

Bumblebee の公式サンプル

https://github.com/elixir-nx/bumblebee/blob/main/notebooks/examples.livemd

実装の全文はこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/bumblebee/question_answering.livemd

## 実行環境

- MacBook Pro 13 inchi
  - 2.4 GHz クアッドコアIntel Core i5
  - 16 GB 2133 MHz LPDDR3
- macOS Ventura 13.0.1
- Rancher Desktop 1.6.2
  - メモリ割り当て 12 GB
  - CPU 割り当て 6 コア

Livebook 0.8.0 の Docker　イメージを元にしたコンテナで動かしました

https://hub.docker.com/layers/livebook/livebook/0.8.0/images/sha256-6ef3fdd439fd72af9ae01d58e947664a46a8493f031e179cc0e91c08ddac02e5?context=explore

コンテナ定義はこちらを参照

https://github.com/RyoWakabayashi/elixir-learning

## セットアップ

必要なモジュールをインストールして EXLA.Backend で Nx が動くようにします

```elixir
Mix.install(
  [
    {:bumblebee, "~> 0.1"},
    {:nx, "~> 0.4"},
    {:exla, "~> 0.4"},
    {:kino, "~> 0.8"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)
```

コンテナで動かしている場合、キャッシュディレクトリーを指定した方が都合がいいです

※詳細は[この記事](https://qiita.com/RyoWakabayashi/items/8070f6bc372015be6a8d)を見てください

```elixir
cache_dir = "/tmp/bumblebee_cache"
```

## モデルのダウンロード

モデルファイルを Haggin Face からダウンロードしてきて読み込みます

必要な場合は `cache_dir` を指定します

```elixir
{:ok, roberta} =
  Bumblebee.load_model({
    :hf,
    "deepset/roberta-base-squad2",
    cache_dir: cache_dir
  })
```

```elixir
{:ok, tokenizer} =
  Bumblebee.load_tokenizer({
    :hf,
    "roberta-base",
    cache_dir: cache_dir
  })
 ```

## 文章の準備

　質問文と、質問の対象になる文章を入力します

```elixir
question_input =
  Kino.Input.text("QUESTION",
    default: "What industries does Elixir help?"
  )
```

![スクリーンショット 2022-12-13 17.25.23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/6a261f9a-af0d-36fc-023a-3ae0bd1bea5b.png)

```elixir
context_input =
  Kino.Input.textarea("CONTEXT",
    default:
      ~s/Elixir is a dynamic, functional language for building scalable and maintainable applications. Elixir runs on the Erlang VM, known for creating low-latency, distributed, and fault-tolerant systems. These capabilities and Elixir tooling allow developers to be productive in several domains, such as web development, embedded software, data pipelines, and multimedia processing, across a wide range of industries./
  )
```

![スクリーンショット 2022-12-13 17.21.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/b8a54389-42f0-2f68-7bdf-c29518b0c460.png)

入力された文章を取得します

```elixir
question = Kino.Input.read(question_input)
context = Kino.Input.read(context_input)
```

## 推論の実行

推論して結果を表示します

この機能は Nx.Serving 用の関数がまだ用意されていないので、前処理、後処理を自分で書きます

```elixir
inputs = Bumblebee.apply_tokenizer(tokenizer, {question, context})
outputs = Axon.predict(roberta.model, roberta.params, inputs)

answer_start_index =
  outputs.start_logits
  |> Nx.argmax()
  |> Nx.to_number()

answer_end_index =
  outputs.end_logits
  |> Nx.argmax()
  |> Nx.to_number()

answer_tokens =
  inputs["input_ids"][[0, answer_start_index..answer_end_index]]
  |> Nx.to_flat_list()

Bumblebee.Tokenizer.decode(tokenizer, answer_tokens)
```

![スクリーンショット 2022-12-13 17.25.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/e4b396e1-8df4-a532-abc3-82c752d5097b.png)

ちゃんと質問に答えました

## まとめ

これからの Bumblebee に期待が持てますね