使用 Milvus Hybrid Search Retriever 实现混合向量检索

技术背景介绍

Milvus 是一个开源的向量数据库，专为嵌入相似度搜索和AI应用而构建。它通过提供统一的用户体验，简化了对非结构化数据的搜索，不论部署环境如何。Milvus 的混合搜索功能将稠密和稀疏向量搜索的优势结合在一起，为复杂搜索任务提供了更好的解决方案。

核心原理解析

Milvus 的混合搜索利用了稠密向量和稀疏向量的独特优势。稠密向量通常由深度学习模型生成，例如 OpenAI 提供的嵌入，而稀疏向量则可以通过经典信息检索算法如 BM25 生成。通过将这两种方法结合，Milvus 可以实现更加准确和高效的搜索结果。

代码实现演示

下面是如何使用 Milvus 实现混合向量检索的完整示例。

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_milvus.retrievers import MilvusCollectionHybridSearchRetriever
from langchain_milvus.utils.sparse import BM25SparseEmbedding
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from pymilvus import (
    Collection,
    CollectionSchema,
    DataType,
    FieldSchema,
    WeightedRanker,
    connections,
)

# 确保已经安装必要的包
# %pip install --upgrade --quiet pymilvus[model] langchain-milvus langchain-openai

# 初始化 Milvus 连接
CONNECTION_URI = "http://localhost:19530"
connections.connect(uri=CONNECTION_URI)

# 定义字段和数据类型
fields = [
    FieldSchema(name="doc_id", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
    FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=1536),
    FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65_535),
]

# 创建 Milvus collection
schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
collection = Collection(name="IntroductionToTheNovels", schema=schema, consistency_level="Strong")

# 定义索引
collection.create_index("dense_vector", {"index_type": "FLAT", "metric_type": "IP"})
collection.create_index("sparse_vector", {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"})
collection.flush()

# 准备示例文本数据
texts = [
    "In 'The Whispering Walls' by Ava Moreno...",
    "In 'The Last Refuge' by Ethan Blackwood...",
    # 省略其他文本以简化展示
]

# 生成稠密和稀疏向量
dense_embedding_func = OpenAIEmbeddings()
sparse_embedding_func = BM25SparseEmbedding(corpus=texts)

# 插入数据
entities = [
    {
        "dense_vector": dense_embedding_func.embed_documents([text])[0],
        "sparse_vector": sparse_embedding_func.embed_documents([text])[0],
        "text": text,
    }
    for text in texts
]
collection.insert(entities)
collection.load()

# 实例化检索器
retriever = MilvusCollectionHybridSearchRetriever(
    collection=collection,
    rerank=WeightedRanker(0.5, 0.5),
    anns_fields=["dense_vector", "sparse_vector"],
    field_embeddings=[dense_embedding_func, sparse_embedding_func],
    field_search_params=[{"metric_type": "IP", "params": {}}, {"metric_type": "IP"}],
    top_k=3,
    text_field="text",
)

# 使用检索器
result = retriever.invoke("What are the stories about ventures?")
for doc in result:
    print(doc.page_content)

# 删除集合
collection.drop()

应用场景分析

Milvus 的混合向量检索在需要处理复杂和多样化数据集的应用中具有显著优势，例如推荐系统、文档搜索和智能问答系统。结合使用稠密和稀疏向量的能力，使其能够处理各种不同类型的查询和数据集。

实践建议

1. 优化索引配置：根据数据集的具体需求和查询模式，调整稠密和稀疏索引的参数设置。
2. 数据预处理：在生成向量表示之前，确保对输入数据进行充分清洗和预处理，以提高搜索精度。
3. 混合重排序策略：利用 WeightedRanker 等机制调整稠密和稀疏向量的权重，以实现更优排序效果。

如果遇到问题欢迎在评论区交流。
—END—