DashVector:グループ化ベクトル検索

背景

シナリオによっては、ベクトル検索の戻り値をグループ化する必要があります。次に例を示します。

• Retrieval-Augmented Generation（RAG）システムでは、ドキュメントは多くの場合複数のセグメントに分割され、各セグメントはベクトル化されてDashVectorに格納されます。多様な結果をベクトル検索で取得するために、ユーザーは異なるドキュメントから最も類似したセグメントが返されることを期待します。

• 商品画像検索のシナリオでは、多くの場合、1 つの商品に対して複数の画像が作成され、各画像はベクトル化されてDashVectorに格納されます。多様な結果をベクトル検索で取得するために、ユーザーは異なる商品の最も類似した画像が返されることを期待します。

DashVector はグループ化ベクトル検索機能をサポートしています。これにより、上記のシナリオで、グループ化ドキュメント検索操作を呼び出し、group_by_field をドキュメント ID と商品 ID にそれぞれ設定できます。詳細については、「グループ化ドキュメント検索」をご参照ください。

例

前提条件

• クラスタが作成されていること。詳細については、「クラスタの作成」をご参照ください。

• API キーを取得していること。詳細については、「API キーの管理」をご参照ください。

• 最新バージョンの SDK がインストールされていること。詳細については、「DashVector SDK のインストール」をご参照ください。

フィールドを持つドキュメントを挿入する

import dashvector
import numpy as np

client = dashvector.Client(
    api_key='YOUR_API_KEY',
    endpoint='YOUR_CLUSTER_ENDPOINT'
)
ret = client.create(
    name='group_by_demo',
    dimension=4,
    fields_schema={'document_id': str, 'chunk_id': int}
)
assert ret

collection = client.get(name='group_by_demo')

ret = collection.insert([
    ('1', np.random.rand(4), {'document_id': 'paper-01', 'chunk_id': 1, 'content': 'xxxA'}),
    ('2', np.random.rand(4), {'document_id': 'paper-01', 'chunk_id': 2, 'content': 'xxxB'}),
    ('3', np.random.rand(4), {'document_id': 'paper-02', 'chunk_id': 1, 'content': 'xxxC'}),
    ('4', np.random.rand(4), {'document_id': 'paper-02', 'chunk_id': 2, 'content': 'xxxD'}),
    ('5', np.random.rand(4), {'document_id': 'paper-02', 'chunk_id': 3, 'content': 'xxxE'}),
    ('6', np.random.rand(4), {'document_id': 'paper-03', 'chunk_id': 1, 'content': 'xxxF'}),
])
assert ret

グループ化ベクトル検索を実行する

ret = collection.query_group_by(
    vector=[0.1, 0.2, 0.3, 0.4],
    group_by_field='document_id',  # document_id フィールドの値で戻り値をグループ化します。
    group_count=2,  # 2 つのグループを返します。
    group_topk=2,   # 各グループから最大 2 つのドキュメントを返します。
)
# 操作が成功したかどうかを確認します。
if ret:
    print('query_group_by success')
    print(len(ret))
    print('------------------------')
    for group in ret:
        print('group key:', group.group_id)
        for doc in group.docs:
            prefix = ' -'
            print(prefix, doc)

サンプル出力は次のとおりです。

query_group_by success
4
------------------------
group key: paper-01
 - {"id": "2", "fields": {"document_id": "paper-01", "chunk_id": 2, "content": "xxxB"}, "score": 0.6807}
 - {"id": "1", "fields": {"document_id": "paper-01", "chunk_id": 1, "content": "xxxA"}, "score": 0.4289}
group key: paper-02
 - {"id": "3", "fields": {"document_id": "paper-02", "chunk_id": 1, "content": "xxxC"}, "score": 0.6553}
 - {"id": "5", "fields": {"document_id": "paper-02", "chunk_id": 3, "content": "xxxE"}, "score": 0.4401}

制限事項