- Retrieval-Augmented Generation(RAG)システムを使うことで、既存の社内知識ベース(Wiki、マニュアル、教育資料、参考資料など)をもとに質問へ回答できるAIアシスタントを作成できる
- PostgreSQL、pgvector、ollama、そして200行以下のGoコードだけでRAGシステムを構築可能
Overview
- 数段落の文章を「ドキュメントコーパス」として使い、各ドキュメントに対してMetaのLlama3を用いてドキュメント埋め込みを生成する(ollamaでローカルホスティング)
- ドキュメントと埋め込みはPostgreSQLのテーブルに保存し、埋め込みの保存とアクセスにはpgvector拡張を利用する
- ユーザークエリに対して最も関連性の高い1件のドキュメントをテーブルから検索し、Llama3を使って応答を生成する
- ollamaはOpenAIに似たHTTP APIを提供し、埋め込みとチャット応答を生成してくれる
- Goコードはjackc/pgxとpgvector-goを使ってPostgresと通信し、ollamaクライアントAPIパッケージを使ってHTTP API呼び出しを処理する
Ollamaでモデルを実行する
- Ollamaはオープンソースモデルをローカルで実行できるようにするツールで、OpenAIスタイルのREST APIを提供する
ollama pull llama3 コマンドでllama3モデルを実行
- ollamaのHTTPサーバーはデフォルトで
127.0.0.1:11434 で利用可能
pgvectorのインストール
- pgvectorはPostgreSQL 12〜16向けの拡張機能で、pgdg APTリポジトリを使っている場合は
sudo apt install postgresql-16-pgvector でインストール可能
- インストール後、
create extension vector; でデータベースに拡張機能を設定
create table items (id serial primary key, doc text, embedding vector(4096)); でドキュメントと埋め込み保存用のテーブルを作成
ドキュメントデータ
- Sherlock Holmesの物語「The Boscombe Valley Mystery」から4つの段落を使用(パブリックドメイン、Project Gutenberg)
コード
- GitHubでMITライセンスとして公開されているデモコードを利用可能
- ドキュメントの挿入は
INSERT INTO items (doc, embedding) VALUES ($1, $2)
- 最も関連性の高いドキュメントの検索は
SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1 (<-> 演算子はpgvectorが提供)
- Ollama APIの呼び出しにはollama Goパッケージを使用
- 埋め込み生成には
api.EmbeddingRequest を使用
- チャット応答の生成には
api.ChatRequest を使用(プロンプトに検索したドキュメントを含める)
コマンドラインインターフェース
ragdemo -insert {path-to-doc-file} でドキュメントをデータベースに保存
ragdemo -query {query-text} でプロンプトを入力して応答を生成
全体の流れ
-insert オプションでドキュメントを保存するとき、ファイル内容を読み込んでLlama3で埋め込みを生成し、PostgreSQLに保存する
-query オプションを使うと、プロンプトの埋め込みを生成した後、items テーブル内の他の埋め込みと比較して「最近傍」のドキュメントを検索する(<-> 演算子でL2距離を計算)
- 検索したドキュメントをプロンプトに含めてLlama3に渡し、チャット応答を生成して出力する
追加のヒント
- 埋め込み生成に特化したモデルの利用を検討する(llama3の代わりに)
- 英語以外の言語については、より適したモデルを探す必要がある
- L2距離以外の距離計算法も試せる(pgvectorは他の方法もサポートしている)
- テーブル全体のスキャンは拡張性が低いため、pgvectorインデックスなどを活用する
- 生成段階でより多くのドキュメントを活用したり、キーワードマッチングなどで追加ドキュメントを取得したりすることも有効
- 生成プロンプトの調整やさまざまなLLMの試行により、出力品質を改善できる
まだコメントはありません。