enterprise-kb-rag/2.3_Milvus向量数据库/05_hybrid_search_rerank.py

import random

from pymilvus import AnnSearchRequest, DataType, MilvusClient, RRFRanker, WeightedRanker

client = MilvusClient(uri="http://localhost:19530")


def demo_rankers():
    # 加权排名：为某些向量场分配更高的权重，从而更加强调这些向量场。
    # 输入值的范围应为 [0,1]，数值越接近 1 表示重要性越高。
    rerank = WeightedRanker(0.8, 0.3)
    print(f'WeightedRanker-->{rerank}')

    # RRFRanker（倒数排序融合）：没有特定重点时使用，k 的默认值为 60，用于控制随着排名增加分数的降低速度。
    ranker = RRFRanker(100)
    print(f'RRFRanker-->{ranker}')


def complex_query():
    # 脚本可重复运行：先清理上一次留下的 collection
    if client.has_collection(collection_name='demo_v3'):
        client.drop_collection(collection_name='demo_v3')

    # # 定义schema
    schema = client.create_schema(enable_dynamic_field=False)
    schema.add_field(field_name='film_id', datatype=DataType.INT64, is_primary=True)
    schema.add_field(field_name='filmVector', datatype=DataType.FLOAT_VECTOR, dim=5)  # 向量字段
    schema.add_field(field_name="posterVector", datatype=DataType.FLOAT_VECTOR, dim=5)  # 向量字段
    # #
    # 定义索引
    index_params = client.prepare_index_params()
    index_params.add_index(field_name='filmVector', index_type="IVF_FLAT",
                           metric_type="L2", params={"nlist": 128})
    # index_type="",  # 留空以使用自动索引 —— 当前 pymilvus/Milvus Lite 版本要求显式指定 index_type，这里用 'AUTOINDEX' 代替课件里的空字符串写法
    index_params.add_index(field_name='posterVector', index_type="AUTOINDEX",
                           metric_type="COSINE")

    # 创建集合
    client.create_collection(collection_name='demo_v3', schema=schema, index_params=index_params)

    # 向量库中插入实体
    entities = []
    for _ in range(1000):
        # 构造实体
        film_id = random.randint(1, 10000)
        film_vector = [random.random() for _ in range(5)]
        poster_vector = [random.random() for _ in range(5)]
        entity = {"film_id": film_id, "filmVector": film_vector, "posterVector": poster_vector}
        entities.append(entity)

    client.insert(collection_name='demo_v3', data=entities)

    # 多向量查询（注意和批量向量查询不同）
    # 多向量搜索使用 hybrid_search() API 在一次调用中执行多个 ANN 搜索请求。每个 AnnSearchRequest 代表特定矢量场上的单个搜索请求。
    # 示例创建两个 AnnSearchRequest 实例以对两个向量字段执行单独的相似性搜索。
    # 创建多搜索请求 filmVector
    query_filmVector = [[0.8896863042430693, 0.370613100114602, 0.23779315077113428, 0.38227915951132996, 0.5997064603128835]]
    dense_search_params = {"data": query_filmVector,
                           "anns_field": "filmVector",  # 该参数值必须与集合模式中使用的值相同。
                           "param": {"metric_type": "L2", "nprobe": 10},  # nprobe代表访问簇的数量
                           "limit": 2}
    request_1 = AnnSearchRequest(**dense_search_params)

    # 创建多搜索请求 posterVector
    query_posterVector = [[0.02550758562349764, 0.006085637357292062, 0.5325251250159071, 0.7676432650114147, 0.5521074424751443]]
    sparse_search_params = {"data": query_posterVector,
                            "anns_field": "posterVector",
                            # 该参数值必须与集合模式中使用的值相同。
                            "param": {"metric_type": "COSINE"},
                            "limit": 2
                            }
    request_2 = AnnSearchRequest(**sparse_search_params)

    reqs = [request_1, request_2]
    ranker = RRFRanker(100)

    res = client.hybrid_search(
        collection_name="demo_v3",
        reqs=reqs,
        ranker=ranker,
        limit=2
    )
    for hits in res:
        print("TopK results:")
        for hit in hits:
            print(hit)


if __name__ == '__main__':
    demo_rankers()
    complex_query()