見出し画像

ベクトルDB構築メモ

Yuki Doi

企業が抱えてるナレッジのベクトル取って、ユーザーの質問にいい感じに答えてくれるDBをある程度汎化して作りたいなー作ってみたぜのメモ

各クラウドインフラにそういう用途のサービス出てきてるんだけど(中の人大変だろうな、、、) 今回はGCPを使ってみる

使った技術

Vertex Matching Engine

python 3.11.3

やりたいこと

インフラ用意して、データとConfig渡すだけでベクトルDBに問い合わせできるライブラリにしたい
あと、Embeddingのレスポンスを、他に正規化されているデータと紐付けたい(ここあんまEmbeddingがどうとか関係ないけど、使いにくいところなので自前実装しないでいいようにしたい)

今回やったこと

概要

Matching Engineへのデータ登録、問い合わせまで
それ以降はまた別記事でお会いしましょう

データの用意

ベクトル化

  • どういうモデルでembeddingするかは自由

  • ここでChatGPT Embeddings APIを使ってみた

  • 元ファイルは普通にGCP Storageに置いて、各環境から読みに行く方針

  • 今回は上記のNetflixのデータのあらすじをEmbeddingした。1000行で5秒ほどの処理時間で意外と速かった。

こんな感じのデータができる (textが元の文章、embeddingがベクトル表現)

                                              text                                          embedding
0  As her father nears the end of his life, filmm...  [0.012684703804552555, 0.001769843278452754, -...
1  After crossing paths at a party, a Cape Town t...  [0.015193070285022259, -0.0153654171153903, -0...
2  To protect his family from a powerful drug lor...  [-0.004565003793686628, -0.008843235671520233,...
3  Feuds, flirtations and toilet talk go down amo...  [0.00742243742570281, -0.017051177099347115, 0...
4  In a city of coaching centers known to train I...  [0.015919573605060577, -0.013753645122051239, ...

ベクトル化したデータの格納

  • ここまででベクトル化されたデータをMatching Engineのファイルフォーマットにする(後述)

  • そのファイルをCloudStorageにアップするとよしなにやってくれる感じ

  • Matching EngineはCloudStorageからembedding済みデータを読む

    • フォーマットは csv, json, avroいずれか

    • json形式の場合、idとembeddingフィールドを持つ

    • そのほかの仕様

  • Indexの更新頻度はバッチとリアルタイム両方ある

Indexの作成

  • Matching EngineのAPIを叩いてIndexを作成する

    • 今回、リアルタイム更新を設定しようとしたがgcloud経由ではできなかったためcurlで直接APIを叩いている。いずれgcloud経由でもできるようになると思われる。

curl -X POST \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  "https://<your-region>-aiplatform.googleapis.com/v1/projects/$(gcloud config get-value project)/locations/<your-region>/indexes" \
  -d @index_metadata.json

curl末尾でindex_metadata.jsonというものを渡しており、これがindexの詳細を記述したもの(どこのデータ使うとかなど)
詳細はこちら
リアルタイム更新にする際はこのjsonのパラメータとしてindexUpdateMethod = STREAM_UPDATEを付け加える。例えばこんな感じ

{
  "display_name": "Search index for netflix shows",
  "metadata": {
    "contentsDeltaUri": "gs://<your-embeddings-bucket>/",
    "config": {
      "dimensions": 1536,
      "approximateNeighborsCount": 100,
      "shardSize": "SHARD_SIZE_SMALL",
      "algorithmConfig": {"treeAhConfig": {}}
    }
  },
  "indexUpdateMethod": "STREAM_UPDATE"
}

  • indexの作成はそこそこ時間がかかるので、Matching EngineのコンソールかAPI叩いてstatusを調べる

    • gcloud ai operations describe <operation-id> --index=<index-id> --project=vector-db-396406 --region=<your-region>

    • operation-idとindex-idはindex作成をAPIコールしたときに表示されるなが〜い数字

    • ファイル内容がビミョ〜に間違ってるとか、Index作成でエラーが出ると結構ツラいと思うよ

Index Endpoint(API)の作成

Matching EngineのAPIを叩いて、上記で作成したIndexを使うAPIを作成する。例えば公開エンドポイント(どこからでもアクセスできるが認証が必要)であればこのように書く

curl -X POST \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  "https://<your-region>-aiplatform.googleapis.com/v1/projects/$(gcloud config get-value project)/locations/<your-region>/indexEndpoints" \
  -d '{"display_name": "public-endpoint-netflix-descriptions", "publicEndpointEnabled": true}'

このIndex EndpointへIndexをデプロイしてようやくAPI完成となる。イメージとしてはAPI(Index Endpoint)へデータ(Index)を繋ぎこむ感じ。

gcloud ai index-endpoints deploy-index INDEX_ENDPOINT_ID \
  --deployed-index-id=DEPLOYED_INDEX_ID \
  --display-name=DEPLOYED_INDEX_NAME \
  --index=INDEX_ID \
  --project=PROJECT_ID \
  --region=LOCATION

Matching Engineへの問い合わせ

やや用意するものが多いが、まあイメージとしては以下の流れ

  1. ユーザーからの問い合わせをベクトルにする

  2. ライブラリ経由でMatching Engine上のAPIを叩いて類似ベクトルを返す

    1. この準備でいろんなインスタンスを作らないといけなくてちょっと面倒

ちなみに2で使うMatchingEngineIndexEndpointのindex_endpoint_nameにはProjectIdからIndex Endpointまでのパスを使う。罠すぎるって笑
こういうやつ projects/<your-project-id>/locations/<your-location>/indexEndpoints/<your-index-endpoint-id>
詳しくはこちら↓

無事に類似ベクトルが取れるとidも一緒に取れる👏 出力例↓

Q: Which action movie should I show that's the best?
A:
s39     distance=0.7712357640266418
s98     distance=0.7672325372695923
s90     distance=0.76556396484375
s29     distance=0.7647834420204163
s49     distance=0.76301109790802
s3      distance=0.7613675594329834
s87     distance=0.7598444223403931
s78     distance=0.7593697309494019
s95     distance=0.7592884302139282
s59     distance=0.7565656900405884

ちなみにs59はNaruto Shippûden the Movie: The Will of Fireで、s3はGanglandsなのでアクションというところはおさえられていそうだった。
これ、やってから思ったけどid列はおそらく重複なければ何入れてもいいので作品タイトルか説明文そのまま入れちゃってもいいかもしれない。その方がわかりやすいしね。
とはいえ、他のデータとの連携を考えると↑みたいなs3とかs59だとかで登録しておくのがいいかな。

感想

Matching Engine、GUIからいじれるところがまだなくて、いちいちAPI叩かないといけない(それもいくつかはcurlで)ため、まだちょっとやりづらい。まあ、一回やってしまえば慣れる程度のものではあるのだけど。
まだ仕方ないけどエラーが出た時の情報も乏しい。まあこれもこなれてくると情報出揃うだろう。
流石にプロダクションコードでpython使う気にはならないのでKotlinあたりで組めるようにライブラリ化しておくけどそれはまた別の記事で。

いいなと思ったら応援しよう！