Google Colaboratory 上の Livebook で LLM (Mistral) の RAG を実装する

## はじめに

前回の記事で Mistral によるテキスト生成を実行しました

https://qiita.com/RyoWakabayashi/items/e2ec30426f554158ef9c

Mistral は日本語を理解してくれるようなので、桃太郎について質問してみましょう

![スクリーンショット 2024-06-01 10.04.46.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/5688eecb-ffc4-5bcf-ce62-a4e8cee2730c.png)

```
質問: 桃太郎の家来の動物は何ですか？
回答: 桃太郎の家来は、猫（ネコ）です。
```

**その桃太郎も読んでみたいけれども**

どうやら Mistral さんは桃太郎を知らないようです

というわけで、 Mistral さんに桃太郎を読んでもらい、その情報を元に回答してもらいます

このように特定の情報（コンテキスト情報 = 文脈）を検索してテキスト生成することを RAG (Retrieval-Augmented Generation) = 検索拡張生成）といいます

本記事では Mistral で RAG を実行します

Bumblebee 公式のノートブックを一部改造して使用しています

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

実装したノートブックはこちら

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

Mistral の RAG について、 Python 実装と詳細解説はこちら

https://docs.mistral.ai/guides/rag/

:::note warn
本記事では RAG 用のテキスト埋め込みに GTE small を使用していますが、日本語には対応していません

日本語に特化した Ruri を使う場合、以下の記事を参照してください

[Livebook 上で日本語テキスト埋め込みモデル Ruri によるテキスト検索を実装する](https://qiita.com/RyoWakabayashi/items/6e12e81606f7ad8e35f6)

[Livebook 上で日本語テキスト埋め込みモデル Ruri によるテキスト検索を実装する（Ollama 利用）](https://qiita.com/RyoWakabayashi/items/463601c2e455d8b0f9f4)
:::

## 大まかな RAG の流れ

- 知識全体をテキストとして読み込む
  知識全体は大きすぎるため LLM に渡せないことが多い
- 知識を一定の大きさでかたまりに分割する
  質問と関係の深いかたまりを検索し、それだけを LLM に渡す
- かたまりを検索するためのインデックスを作成する
  - 各かたまりのテキストを特徴量ベクトルに変換する（= テキスト埋め込み）
  - 特徴量ベクトルをインデックスに登録する
- 質問のテキストも特徴量ベクトルに変換し、インデックスから関係の深いかたまりを検索する
- 検索結果の上位を繋げてコンテキスト情報（文脈情報）にする
- コンテキスト情報を踏まえて質問に回答するよう LLM に指示する

![rag.drawio.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2cbcc650-ea9a-2117-42bc-4d027580148d.png)

## 事前準備

### 各種サービス、モデルの準備

前回の記事と同様、 ngrok 、 HuggingFace 、 Google Colaboratory 有償版を使います

Mistral モデルの使用条件にも合意しておきましょう

https://qiita.com/RyoWakabayashi/items/e2ec30426f554158ef9c#%E4%BA%8B%E5%89%8D%E6%BA%96%E5%82%99

### コンテキスト情報の準備

今回のコンテキスト情報は桃太郎です

青空文庫から桃太郎のテキストを取得し、平文テキストとして GitHub にアップロードしておきます

https://www.aozora.gr.jp/

今回は楠山正雄さんの書いた桃太郎（著作権切れ）を使用します

https://www.aozora.gr.jp/cards/000329/files/18376_12100.html

青空文庫の収録ファイル取り扱い基準についてはこちら

https://www.aozora.gr.jp/guide/kijyunn.html

ルビは邪魔なので削除します

http://www.kepe.net/ruby/

ルビを削除した文書をテキストファイルとして保存し、 GitHub リポジトリーにコミット、プッシュします

準備したものがこちらです

https://raw.githubusercontent.com/RyoWakabayashi/elixir-learning/main/livebooks/bumblebee/colab/momotaro.txt

## Livebook の起動

以下の記事を参考に Google Colaboratory 上で Livebook を起動しましょう

ただし、ランタイムには L4 GPU を指定してください

https://qiita.com/RyoWakabayashi/items/113b94866780c7646af1

![スクリーンショット 2024-05-30 22.15.28.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/a26c2e10-0e4a-846e-836a-c14a34b41ed5.png)

T4 GPU ではメモリ不足で実行できません

実装したノートブックはこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/colab/Livebook.ipynb


## ノートブックの準備

新しいノートブックを開きます

左メニューから南京錠のアイコンをクリックして「SECRETS」を表示し、「+ New secret」をクリックします

![スクリーンショット 2024-05-30 21.51.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/077f902b-126b-c285-e67b-30cd2786fb72.png)

「Name」を「HF_TOKEN」、 Value を Hugging Face のトークンにして「+ Add」をクリックします

![スクリーンショット 2024-05-30 21.52.09.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/c65ba1aa-98b6-92c0-8717-07812438aa8d.png)

## 依存モジュールのインストール

以下のコードを実行します

```elixir
Mix.install(
  [
    {:bumblebee, "~> 0.5"},
    {:nx, "~> 0.7"},
    {:exla, "~> 0.7"},
    {:kino, "~> 0.12"},
    {:hnswlib, "~> 0.1"},
    {:req, "~> 0.4"}
  ],
  system_env: [
    {"XLA_TARGET", "cuda120"},
    {"EXLA_TARGET", "cuda"}
  ]
)

Nx.global_default_backend(EXLA.Backend)
```

[Bumblebee](https://github.com/elixir-nx/bumblebee)、 [Nx](https://github.com/elixir-nx/nx/tree/main/nx)、 [EXLA](https://github.com/elixir-nx/nx/tree/main/exla)、 [Kino](https://github.com/livebook-dev/kino)、 [HNSWLib](https://github.com/elixir-nx/hnswlib)
、 [Req](https://github.com/wojtekmach/req) の 2024年6月1日時点の最新版をインストールしています

EXLA　で GPU を使うため、環境変数の設定もしています

- XLA_TARGET: cuda120
- EXLA_TARGET: cuda

2024年5月30日時点では Google Colaboratory で CUDA 12.2 が使えるため、 `cuda120` を指定しています

各バージョン指定はそのときの最新を使うなど、適宜変更してください

## 知識の準備

Mistral に与える知識を準備します

用意しておいた桃太郎のテキストファイルを Req で読み込みます

```elixir
%{body: text} =
  Req.get!(
    "https://raw.githubusercontent.com/RyoWakabayashi/elixir-learning/main/livebooks/bumblebee/colab/momotaro.txt"
  )

IO.puts("Document length: #{String.length(text)}")
```

実行結果
```elixir
Document length: 5863
```

テキストを 1024 文字毎のかたまりに分割します

```elixir
chunks =
  text
  |> String.codepoints()
  |> Enum.chunk_every(1024)
  |> Enum.map(&Enum.join/1)

length(chunks)
````

今回の場合は 6 個のかたまりになります

テキスト埋め込みのモデルとして、 GTE-small を読み込みます

https://huggingface.co/thenlper/gte-small

```elixir
repo = {:hf, "thenlper/gte-small"}

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

:ok
```

テキスト埋め込みを子プロセスとして起動します

```elixir
serving =
  Bumblebee.Text.TextEmbedding.text_embedding(model_info, tokenizer,
    compile: [batch_size: 64, sequence_length: 512],
    defn_options: [compiler: EXLA],
    output_attribute: :hidden_state,
    output_pool: :mean_pooling
  )

Kino.start_child({Nx.Serving, serving: serving, name: GteServing})
```

テキスト埋め込みを実行します

```elixir
results = Nx.Serving.batched_run(GteServing, chunks)
chunk_embeddings = for result <- results, do: result.embedding

List.first(chunk_embeddings)
```

実行結果
```elixir
#Nx.Tensor<
  f32[384]
  [-0.3020172715187073, -0.4361937642097473, 0.5629637241363525, -0.21154245734214783, ...]
>
```

テキストを特徴量ベクトルに変換できました

## インデックスの作成と検索

知識から回答に使用するかたまりを検索するためのインデックスを作成します

```elixir
{:ok, index} = HNSWLib.Index.new(:cosine, 384, 1_000_000)

for embedding <- chunk_embeddings do
  HNSWLib.Index.add_items(index, embedding)
end

HNSWLib.Index.get_current_count(index)
```

質問を定義し、質問自体も特徴量ベクトルに変換します

インデックスを使用し、質問に関係する知識のかたまりを取得します

```elixir
query = "桃太郎の家来の動物は何ですか？"

%{embedding: embedding} = Nx.Serving.batched_run(GteServing, query)

{:ok, labels, dist} = HNSWLib.Index.knn_query(index, embedding, k: 4)
```

実行結果
```elixir
{:ok,
 #Nx.Tensor<
   u64[1][4]
   EXLA.Backend<cuda:0, 0.2360341075.719716378.206135>
   [
     [4, 3, 1, 5]
   ]
 >,
 #Nx.Tensor<
   f32[1][4]
   EXLA.Backend<cuda:0, 0.2360341075.719716378.206137>
   [
     [0.08875870704650879, 0.09236758947372437, 0.11106133460998535, 0.11499273777008057]
   ]
 >}
```

検索結果の番号から、コンテキスト情報として使うテキストを取得します

```elixir
# We can see some overlapping in our chunks
context =
  labels
  |> Nx.to_flat_list()
  |> Enum.sort()
  |> Enum.map(fn idx -> "[...] " <> Enum.at(chunks, idx) <> " [...]" end)
  |> Enum.join("\n\n")

IO.puts(context)
```

## 回答の生成

テキスト生成のために Mistral のモデルを読み込みます

```elixir
hf_token = System.fetch_env!("LB_HF_TOKEN")
repo = {:hf, "mistralai/Mistral-7B-Instruct-v0.2", auth_token: hf_token}

{:ok, model_info} = Bumblebee.load_model(repo, type: :bf16)
{:ok, tokenizer} = Bumblebee.load_tokenizer(repo)
{:ok, generation_config} = Bumblebee.load_generation_config(repo)

generation_config = Bumblebee.configure(generation_config, max_new_tokens: 100)

:ok
```

テキスト生成を子プロセスとして起動します

```elixir
serving =
  Bumblebee.Text.generation(model_info, tokenizer, generation_config,
    compile: [batch_size: 1, sequence_length: 6000],
    defn_options: [compiler: EXLA]
  )

Kino.start_child({Nx.Serving, name: MistralServing, serving: serving})
```

テキスト生成プロセスにコンテキスト情報と質問を与え、回答を取得します

```elixir
prompt =
  """
  コンテキスト情報は以下の通りです.
  ---------------------
  #{context}
  ---------------------
  与えられたコンテキスト情報に基づき、事前の知識なしに質問に答えてください.
  質問: #{query}
  回答:
  """

results = Nx.Serving.batched_run(MistralServing, prompt)
```

実行結果
```elixir
%{
  results: [
    %{
      text: "桃太郎の家来は、犬、猿、きじです。",
      token_summary: %{input: 4205, output: 22, padding: 1795}
    }
  ]
}
```

![スクリーンショット 2024-06-01 10.49.27.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/ad576110-91d9-e40a-e986-4b8f680b88a3.png)

桃太郎の家来を正しく答えてくれました

本文中の「きじ」が平仮名になっているため、回答でも平仮名で返っています

## まとめ

Livebook 上で LLM の RAG が実行できました

LLM で出来ることの幅が広がりますね