分组向量检索

背景介绍

在向量检索的实际应用中，有些场景需要将向量检索的结果分组返回。例如：

• 在RAG中，一篇文档往往需要拆分为多个段落，每个段落生成一个向量存入DashVector。在向量检索时，为了结果的多样性，不希望所有结果都来自同一篇文档的段落，而是希望结果返回多篇文档，并且每篇文档下仅返回最相似的若干个段落。

• 在商品图像检索时，每个商品通常有多个商品图片，每个图片生成一个向量存入DashVector。在向量检索时，为了结果的多样性，不希望所有结果都是同一个商品的图片，而是希望返回多样化商品，并且每个商品下仅返回最相似的若干个图片。

向量检索服务DashVector支持分组向量检索，对于上面的两个场景可以通过分组检索Doc接口分别设置group_by_field为“文档ID”和“商品ID”，然后执行分组向量检索。

使用示例

前提条件

• 已创建Cluster：创建Cluster。

• 已获得API-KEY：API-KEY管理。

• 已安装最新版SDK：安装DashVector SDK。

插入带有Field的数据

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}

限制说明