タイトル：純Goの機械学習フレームワーク「tensai」でRAGを作ってみる

以前、純Goの機械学習フレームワーク「tensai」でXOR学習とローカルLLMを試すという記事を書きました。

今回は、その続きとしてtensaiを利用した簡単なRAGを作ってみます。

RAGとは

RAGは、質問に関連する文書を検索し、その検索結果をLLMへ渡して回答を生成する仕組みです。

今回は以下の流れで実装しました。

1. docs/内のMarkdownとテキストファイルを読み込む
2. 文書をチャンクに分割する
3. 文書と質問をTF-IDFベクトルへ変換する
4. コサイン類似度が高いチャンクを検索する
5. 検索結果をtensaiのOpenAI互換APIへ渡す
6. ローカルLLMで回答を生成する

サンプルプログラム

今回作成したコードはこちらです。

GitHub – taako-502/tensai-rag-example

Go 1.27以上と、回答生成に利用するモデルが必要です。

検索だけであれば、モデルを起動しなくても試せます。

go run ./cmd/rag \
  -docs ./docs \
  -retrieve-only \
  "RAGはどうやって文書を検索しますか？"

実行すると、質問と関連性の高い文書が順番に表示されます。

indexed 2 chunks from 2 documents
[1] rag.md (score: 0.4519)
[2] tensai.md (score: 0.0737)

スコアは確率ではなく、質問と文書のコサイン類似度です。値が大きいほど、質問との関連性が高いと判断できます。

ローカルLLMで回答を生成する

最初にtensaiサーバーを起動します。

go run github.com/mattn/tensai/cmd/tensai@v0.0.26 serve -addr 127.0.0.1:8080

別のターミナルでRAGを実行します。

go run ./cmd/rag -docs ./docs

質問を入力すると、関連文書を検索したあと、その内容をローカルLLMへ渡して回答を生成します。

% go run ./cmd/rag -docs ./docs
indexed 2 chunks from 2 documents
質問: tensaiとは？
[1] source: tensai.md

検索結果:
[1] tensai.md (score: 0.1646)
[2] rag.md (score: 0.0472)

単に検索結果を整形しているのではなく、検索した文書を根拠としてLLMに回答を生成させています。

今回学んだこと

品質の高いRAGを作るためには、大きく3つのことを考える必要があると分かりました。

1. 最初に読み込ませる文書をどうするか

LLMが回答するときの根拠は、最初に用意した文書です。

必要な情報が文書に含まれていなければ、検索やプロンプトを改善しても正しい回答は生成できません。

そのため、どの文書を用意するか、どの単位で分割するかを考える必要があります。

2. LLMへ渡すプロンプトをどうするか

検索結果を取得したあとは、その内容と質問をLLMへ渡します。

今回のサンプルでは、次のような指示を入れています。

• 検索結果だけを根拠に回答する
• 根拠がない場合は、情報が見つからないと回答する
• 回答に参照した文書番号を付ける

検索結果が同じでも、プロンプトによって回答の内容や形式が変わります。

3. 検索結果の品質をどう評価するか

期待する文書が検索結果の上位に表示されているかを確認する必要があります。

今回のサンプルでは、-retrieve-onlyを使うとLLMを呼び出さず、検索結果と類似度スコアだけを確認できます。

go run ./cmd/rag -docs ./docs -retrieve-only "tensaiとは？"

質問ごとに期待する文書を用意し、その文書が上位に取得できるかをテストすると検索品質を評価しやすくなります。

LLMへ渡す前の検索結果が間違っている場合、プロンプトだけを改善しても正しい回答にはなりません。

制約

今回の検索処理は、意味埋め込みではなくTF-IDFを利用しています。

単語や文字の一致をもとに検索するため、「自動車」と「車」のように意味は近くても表現が異なる文章の検索は苦手です。

また、ベクトルは保存せず、起動するたびにメモリ上で作り直しています。

本格的に利用する場合は、埋め込みモデル、ベクトルデータベース、最低スコア、メタデータによる絞り込みなども検討する必要があります。

付録：用語集

まとめ

tensaiを使って、文書の検索からローカルLLMによる回答生成まで、RAGの基本的な流れを確認できました。

実際に作ってみると、品質の高いRAGにするためには、次の3つが重要だと分かりました。

1. 最初に読み込ませる文書をどうするか
2. LLMに回答させるためのプロンプトをどうするか
3. 検索結果の品質をどう評価するか

RAGはLLMへ文書を渡すだけではなく、その前段にある文書と検索の設計も重要だと感じました。