Vertex AI Matching Engineでフィルタ機能を使う

ベクトルのフィルタ機能

前回の記事にて、Matching Engineを触ってみました。

https://zenn.dev/articles/2bdde4e7169eed

実際に使う場合は、特定のデータにフィルタした上で、近傍探索をしたい場合があると思います。

例えば、公開・非公開状態の記事がある中で、公開されている記事のみでベクトル検索を行い関連する記事のデータを取得したい場合などです。事前のフィルタができないと、ベクトル検索結果後にフィルタされ十分な検索結果が残らなかったりする可能性があります。

ElasticSearchは事前のフィルタ機能があります。faissは記事執筆時点の2023/05月だとそういった機能はありません。

今回は、Matching Engineのフィルタ機能を使ってみようと思います。

データの作成

前記事のデータ作成の部分を流用し、"restricts": [{"namespace": "class","allow": [str(i%10)]}]}という項目を追加しています。namespaceがグループで、allow以下にグループの値を持たせることができます。

namespaceは複数定義してもOK。allowも複数の値があっても問題ないみたいです。

import json
import random

# JSONファイルを開く
with open('index_data.json', 'w') as f:
    # 改行区切りのJSONオブジェクトを含む文字列を生成する
    data_str = ''
    for i in range(100):
        obj = {'id': f'{i}',
               'embedding': [round(random.random(),5) for _ in range(100)],
+              'restricts': [{"namespace": "class","allow": [str(i%10)]}]
                }
        obj_str = json.dumps(obj)
        data_str += obj_str + '\n'

    # 文字列をファイルに書き込む
    f.write(data_str)

とりあえず、下記のような形でデータを定義しました。

インデックス作成/エンドポイント作成/デプロイ

割愛します。前回の記事を参考にしてみてください。

検索

下記で特定のグループに絞った上で、検索を行なっています。

import random

request = aiplatform_v1beta1.FindNeighborsRequest(
   index_endpoint="projects/プロジェクト/locations/ロケーション/indexEndpoints/INDEX_ENDPOINT_ID",
   deployed_index_id="DEPLOYED_INDEX_ID",
)
dp1 = aiplatform_v1beta1.IndexDatapoint(
   datapoint_id="0",
   feature_vector=[round(random.random(), 5) for _ in range(100)],
+  restricts=[{"namespace": "class", "allow_list": ["3"]}]
)
query = aiplatform_v1beta1.FindNeighborsRequest.Query(
   datapoint=dp1,
)
request.queries.append(query)

response = client.find_neighbors(request)

フィルタはidが3で割り切れるやつのみにフィルタしてから、ベクトル検索を行っています。

最後に

今回の記事、情報量は少ないですが、前の記事と合わせるとごちゃごちゃするので分離しました。

フィルタ機能、使い道が多そうで良い機能だと思いました。