🐥

GraphRAG を TiDB Serverless Vector Search で起動する

2024/11/14に公開

RAG

tidbserverless

GraphRAG

tech

今日は最近注目を集めているGraphRAGをTiDB ServerlessでVector Search 起動します。なおGraphRAGは現在私も勉強中なので作ることはまだできません。PingCAPがTiDB.AIとしてサンプルデモを提供してくれているので、まずは学習用基盤としてそれを起動する手順を纏めておきます。

GraphRAGとは

過去以下の記事でVector Searchを用いたRAGを作る手順をまとめました。

要はある特定のドキュメントをベクトル化してVector Searchに格納しておき、自然言語で検索を掛けます。その際自然言語をベクトルデータに変換したり、回答を自然言語としてスムーズなものに要約するために外部LLMを利用します。

GraphRAGとは、RAGの弱点を補完する新しい技術として注目を集めており、グラフ理論に基づいたデータの構造化を行いそこから関係性を整理しています。

こういう感じで各頂点（ノード）にデータの属性を保持し、ノード間を結ぶ線でその関係性を維持します。

従来のRAGは対象データをチャンクという小さいサイズに分割してベクトル化し検索対象とします。その際、チャンクが小さすぎると文節が細切れとなり、各言葉の相関が失われてしまいます。一方チャンクを大きくすると重要な用語とまとめて他の言葉がベクトル化されるため、検索結果のノイズが大きくなります。
このため言葉の管理対象となるべく小さくしつつ、その相関性を維持しやすくすることを期待されているのがGraphRAGと言われるものであり精度の向上が期待されています。（必ず向上するわけでは勿論ありません）
日本オラクル社の@ksonoda さんの記事がとてもわかりやすかったので共有させていただきます。是非いいねをお願いします。

TiDB Serverless Vector Search との関係性

過去のVector Searchを用いたRAGの記事では、検索対象の文書をチャンクという単位にぶつ切りにして、それをベクトル化することで検索を実現していました。GraphRAGにおいても、データの整理を行うための構造や、各ノード間を指し示す矢印の情報はベクトル化され検索対象になります。したがって基礎技術は共通しています。ベクトルは引き続き使われるが、データの持ち方が異なる、というイメージです。
また世の中にはグラフデータの取り扱いに特化したデータベースも存在しています。代表的なものでいえばNeo4jやAmazon Neptuneなどがありますが、TiDBはグラフデータベースを提供しているわけではありません。あくまでGraphRAGの中核をなすベクトルデータを取りあつかうベクトルデータベースを提供しています。
TiDB.AIはPingCAPが提供しているため全てがTiDB Serverlessで構成されていますが、例えば各ノードの情報、各矢印が持つ属性情報は（その矢印にはどういう関係が含まれているか？）はベクトル化させベクトルデータベースを用いて、各ノード間の矢印情報（だれがだれと関係を保有している）はグラフデーベースを用いることで大規模環境ではクエリの高速化が期待できる（かもしません。すいません、まだ勉強中なのでこのレベルの言い方で申し訳ありません）。

さっそくやってみる

https://tidb.ai/docs
に従って作業を行っていきます。環境はAmazon Linux 2023, LLMモデルはOpenAIを使いますが、基本コンテナが動作すればどこでも動作しますし、LLMモデルも複数から選択できます。EBSストレージではデフォルトの8GBではなく20GBにしておいてください。

1. EC2 インスタンスの起動とDocker/Docker Compose のインストール

提供されているコンテナが大きいのでt2.xlargeを選択します。
起動したら以下をインストールします。

sudo yum update -y
sudo amazon-linux-extras enable docker
sudo yum install -y docker
sudo systemctl start docker
sudo systemctl enable docker

Dockerのインストールが終わったら次にDocker Composeです

VERSION="v2.20.0"
sudo curl -L https://github.com/docker/compose/releases/download/${VERSION}/docker-compose-$(uname -s)-$(uname -m) -o /usr/local/bin/docker-compose
sudo chmod +x /usr/local/bin/docker-compose

2.gitのインストールとレポジトリのクローン

sudo yum install git

レポジトリをクローンします。

git clone https://github.com/pingcap/tidb.ai.git
cd tidb.ai

3.OpenAI APIのAPI Keyを入手

OpenAIのAPI Keyはブラウザから利用する物とは別に購入する必要があるので注意してください。詳細な手順はこの記事では割愛しますが以下から入手可能です。

私の環境では大体１回のクエリに0.05＄消費していました。skから始まる文字列になります。

4.Python3 を用いたデモアプリ用シークレットキーの作成

Amazon Linux 2023 ではあらかじめPython3が入っていますのでインストールは不要です。

python3 -c "import secrets; print(secrets.token_urlsafe(32))"

生成された文字列をコピーしておきます。

5.環境設定ファイルの修正

.envファイルをエディタで開きます。

cp .env.example .env
vi .env

SECRET_KEY,TIDB_HOST, TIDB_USER, TIDB_PASSWORD,TIDB_DATABASEに必要な値を書き込み保存します。

6.DBスキーマの構築とブートストラップ

以下を実行しDBにスキーマを構築します。

docker-compose run backend /bin/sh -c "alembic upgrade head"

実行後TiDB Cloud側で見るとTableなどが出来ていることがわかります。

次にブートストラップとしてデータの書き込みやアプリ用ユーザーの作成を行います。

docker-compose run backend /bin/sh -c "python bootstrap.py"

画面に以下が表示されます。これは後ほど使いますのでコピーをしておいてください。

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
[IMPORTANT] Admin user created with email: admin@example.com and password: ZHwDj2xxxxxxxxxxxxxxxx
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

7. 起動とアプリへのアクセス

ではいよいよ起動です。

docker-compose --profile local-embedding-reranker up

起動が完了したらWebインスタンスが起動していますのでhttp://<ip address>:3000 へアクセスします。

8. ログインと初期設定

先ほどコピーしておいたID/Passwordを入力します。
次にOpenAIを選択しAPI KeyをセットしCrete LLMボタンをクリックします。

次にEmbedding（文字列をベクトルに変換する作業）に使うOpenAIをセットしCreate Embeddingをクリックします。

最後にチャットの土台となる情報を入力します。ファイル、単独WEBページ、特定WEBディレクトリから　選択できますがこの例では以下を用います。

Create DatasourceをクリックしたらOKをクリックして利用開始です。