分組向量檢索

DashVector發表於2024-11-25

本文介紹如何在向量檢索時將結果按照欄位值進行分組返回。

背景介紹

在向量檢索的實際應用中，有些場景需要將向量檢索的結果分組返回。例如：

  • 在RAG中，一篇文件往往需要拆分為多個段落，每個段落生成一個向量存入DashVector。在向量檢索時，為了結果的多樣性，不希望所有結果都來自同一篇文件的段落，而是希望結果返回多篇文件，並且每篇文件下僅返回最相似的若干個段落。

  • 在商品影像檢索時，每個商品通常有多個商品圖片，每個圖片生成一個向量存入DashVector。在向量檢索時，為了結果的多樣性，不希望所有結果都是同一個商品的圖片，而是希望返回多樣化商品，並且每個商品下僅返回最相似的若干個圖片。

向量檢索服務DashVector支援分組向量檢索，對於上面的兩個場景可以透過分組檢索Doc介面分別設定group_by_field為"文件ID"和"商品ID"，然後執行分組向量檢索。

使用示例

前提條件

  • 已建立Cluster
  • 已獲得API-KEY
  • 已安裝最新版SDK

插入帶有Field的資料

說明

需要使用您的api-key替換示例中的 YOUR_API_KEY、您的Cluster Endpoint替換示例中的YOUR_CLUSTER_ENDPOINT，程式碼才能正常執行。

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個doc
)
# 判斷是否成功
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}

限制說明

重要

  1. group_by_field只能指定新建Collection時透過fields_schema引數定義的Field名稱，Schema Free欄位不支援分組檢索。

  2. group_count和group_topk均為盡力而為引數，實際返回的分組數（group_count）和每個分組的doc數（group_topk）可能少於所設定的值。DashVector會優先保證分組數（group_count）。

  3. 過大的group_count和group_topk會增加索引掃描量，從而導致介面耗時增加。當前group_count最大值為64，group_topk最大值為16。

相關文章