使用 Milvus 2.5 开始混合语义/全文搜索

Engineering

December 17, 2024

Stefan Webb

在本文中，我们将向您展示如何快速使用新的全文搜索功能，并将其与基于向量嵌入的传统语义搜索相结合。

要求

首先，确保已安装 Milvus 2.5：

pip install -U pymilvus[model]

并使用Milvus 文档中的安装说明运行 Milvus Standalone 实例（例如在本地计算机上）。

构建数据 Schema 和搜索索引

我们导入所需的类和函数：

from pymilvus import MilvusClient, DataType, Function, FunctionType, model

你可能已经注意到 Milvus 2.5 的两个新条目：Function 和FunctionType ，我们稍后将对此进行解释。

接下来，我们用 Milvus Standalone（即本地）打开数据库，并创建数据模式。该模式包括一个整数主键、一个文本字符串、一个维数为 384 的稠密向量和一个稀疏向量（维数不受限制）。请注意，Milvus Lite 目前不支持全文检索，只支持 Milvus Standalone 和 Milvus Distributed。

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

schema = client.create_schema()

schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1000, enable_analyzer=True)
schema.add_field(field_name="dense", datatype=DataType.FLOAT_VECTOR, dim=768),
schema.add_field(field_name="sparse", datatype=DataType.SPARSE_FLOAT_VECTOR)

{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'dense', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 768}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>}], 'enable_dynamic_field': False}

你可能已经注意到enable_analyzer=True 参数。这将告诉 Milvus 2.5 在这个字段上启用词法分析器，并建立一个词组和词组频率列表，这是全文搜索所必需的。sparse 字段将保存文档的向量表示，即由解析产生的词包text 。

但我们如何连接text 和sparse 字段，并告诉 Milvussparse 应如何从text 计算出来？这就需要调用Function 对象并将其添加到 Schema 中：

bm25_function = Function(
    name="text_bm25_emb", # Function name
    input_field_names=["text"], # Name of the VARCHAR field containing raw text data
    output_field_names=["sparse"], # Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
    function_type=FunctionType.BM25,
)

schema.add_function(bm25_function)

{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'dense', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 768}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>, 'is_function_output': True}], 'enable_dynamic_field': False, 'functions': [{'name': 'text_bm25_emb', 'description': '', 'type': <FunctionType.BM25: 1>, 'input_field_names': ['text'], 'output_field_names': ['sparse'], 'params': {}}]}

Function 对象的抽象比应用全文检索更通用。将来，它可能会用于一个字段需要成为另一个字段的函数的其他情况。在我们的案例中，我们通过函数FunctionType.BM25 指定sparse 是text 的函数。BM25 指的是信息检索中的一个常用度量，用于计算查询与文档（相对于文档 Collections）的相似度。

我们使用 Milvus 中的默认嵌入模型，即paraphrase-albert-small-v2：

embedding_fn = model.DefaultEmbeddingFunction()

下一步是添加搜索索引。我们有一个用于密集向量，另一个用于稀疏向量。索引类型为SPARSE_INVERTED_INDEX ，BM25 ，因为全文搜索所需的搜索方法与标准密集向量的搜索方法不同。

index_params = client.prepare_index_params()

index_params.add_index(
    field_name="dense",
    index_type="AUTOINDEX", 
    metric_type="COSINE"
)

index_params.add_index(
    field_name="sparse",
    index_type="SPARSE_INVERTED_INDEX", 
    metric_type="BM25"
)

最后，我们创建我们的 Collections：

client.drop_collection('demo')
client.list_collections()

[]

client.create_collection(
    collection_name='demo', 
    schema=schema, 
    index_params=index_params
)

client.list_collections()

['demo']

这样，我们就有了一个空数据库，可以接受文本文档并执行语义和全文检索！

插入数据和执行全文搜索

插入数据与 Milvus 以前的版本没有什么不同：

docs = [
    'information retrieval is a field of study.',
    'information retrieval focuses on finding relevant information in large datasets.',
    'data mining and information retrieval overlap in research.'
]

embeddings = embedding_fn(docs)

client.insert('demo', [
    {'text': doc, 'dense': vec} for doc, vec in zip(docs, embeddings)
])

{'insert_count': 3, 'ids': [454387371651630485, 454387371651630486, 454387371651630487], 'cost': 0}

在进行混合搜索之前，让我们先演示一下全文搜索：

search_params = {
    'params': {'drop_ratio_search': 0.2},
}

results = client.search(
    collection_name='demo', 
    data=['whats the focus of information retrieval?'],
    output_fields=['text'],
    anns_field='sparse',
    limit=3,
    search_params=search_params
)

搜索参数drop_ratio_search 指的是在搜索算法中放弃得分较低文档的比例。

让我们来看看结果：

for hit in results[0]:
    print(hit)

{'id': 454387371651630485, 'distance': 1.3352930545806885, 'entity': {'text': 'information retrieval is a field of study.'}}
{'id': 454387371651630486, 'distance': 0.29726022481918335, 'entity': {'text': 'information retrieval focuses on finding relevant information in large datasets.'}}
{'id': 454387371651630487, 'distance': 0.2715056240558624, 'entity': {'text': 'data mining and information retrieval overlap in research.'}}

执行语义和全文混合搜索

现在，让我们结合所学知识来执行混合搜索，将单独的语义搜索和全文搜索与 Reranker 结合起来：

from pymilvus import AnnSearchRequest, RRFRanker
query = 'whats the focus of information retrieval?'
query_dense_vector = embedding_fn([query])

search_param_1 = {
    "data": query_dense_vector,
    "anns_field": "dense",
    "param": {
        "metric_type": "COSINE",
    },
    "limit": 3
}
request_1 = AnnSearchRequest(**search_param_1)

search_param_2 = {
    "data": [query],
    "anns_field": "sparse",
    "param": {
        "metric_type": "BM25",
        "params": {"drop_ratio_build": 0.0}
    },
    "limit": 3
}
request_2 = AnnSearchRequest(**search_param_2)

reqs = [request_1, request_2]

ranker = RRFRanker()

res = client.hybrid_search(
    collection_name="demo",
    output_fields=['text'],
    reqs=reqs,
    ranker=ranker,
    limit=3
)
for hit in res[0]:
    print(hit)

{'id': 454387371651630485, 'distance': 0.032786883413791656, 'entity': {'text': 'information retrieval is a field of study.'}}
{'id': 454387371651630486, 'distance': 0.032258063554763794, 'entity': {'text': 'information retrieval focuses on finding relevant information in large datasets.'}}
{'id': 454387371651630487, 'distance': 0.0317460335791111, 'entity': {'text': 'data mining and information retrieval overlap in research.'}}