![見出し画像](https://assets.st-note.com/production/uploads/images/113967299/rectangle_large_type_2_2c6146e7c3844389c970ec21e1b2512a.png?width=1200)
ベクトルDB構築メモ
企業が抱えてるナレッジのベクトル取って、ユーザーの質問にいい感じに答えてくれるDBをある程度汎化して作りたいなー作ってみたぜのメモ
各クラウドインフラにそういう用途のサービス出てきてるんだけど(中の人大変だろうな、、、) 今回はGCPを使ってみる
使った技術
Vertex Matching Engine
python 3.11.3
やりたいこと
インフラ用意して、データとConfig渡すだけでベクトルDBに問い合わせできるライブラリにしたい
あと、Embeddingのレスポンスを、他に正規化されているデータと紐付けたい(ここあんまEmbeddingがどうとか関係ないけど、使いにくいところなので自前実装しないでいいようにしたい)
今回やったこと
概要
Matching Engineへのデータ登録、問い合わせまで
それ以降はまた別記事でお会いしましょう
データの用意
データはNetflixの作品を監督やキャスト、ジャンル、あらすじに分けたcsvデータを使った
https://www.kaggle.com/datasets/shivamb/netflix-shows?resource=download
将来的にGCP上だけで自動Embedding→Matching Engine自動更新をやるかもしれないのでこの元データもCloud Storageに配置した
ベクトル化
どういうモデルで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への問い合わせ
やや用意するものが多いが、まあイメージとしては以下の流れ
ユーザーからの問い合わせをベクトルにする
ライブラリ経由でMatching Engine上のAPIを叩いて類似ベクトルを返す
この準備でいろんなインスタンスを作らないといけなくてちょっと面倒
ちなみに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あたりで組めるようにライブラリ化しておくけどそれはまた別の記事で。
この記事が気に入ったらサポートをしてみませんか?