教程：创建混合检索管道

在 Colab 中打开下载

_{最后更新：2025 年 8 月 25 日}

级别：中级
完成时间：15 分钟
使用的组件：DocumentSplitter、SentenceTransformersDocumentEmbedder、DocumentJoiner、InMemoryDocumentStore、InMemoryBM25Retriever、InMemoryEmbeddingRetriever 和 TransformersSimilarityRanker
先决条件：无
目标：完成本教程后，您将了解如何创建混合检索以及它何时有用。

概述

混合检索结合了基于关键词的检索和基于嵌入的检索技术，利用了这两种方法的优点。本质上，密集嵌入在理解查询的上下文细微差别方面表现出色，而基于关键词的方法在匹配关键词方面表现出色。

在许多情况下，简单的基于关键词的方法（如 BM25）比密集检索效果更好（例如，在医疗保健等特定领域），因为密集模型需要针对数据进行训练。有关混合检索的更多详细信息，请参阅博客文章：混合文档检索。

准备 Colab 环境

安装 Haystack

使用 pip 安装 Haystack 和其他必需的包

%%bash

pip install haystack-ai
pip install "datasets>=2.6.1"
pip install "sentence-transformers>=4.1.0"
pip install accelerate

初始化 DocumentStore

您将通过初始化一个 DocumentStore 来开始创建您的问答系统。DocumentStore 存储您的系统用于查找问题答案的文档。在本教程中，您将使用 InMemoryDocumentStore。

from haystack.document_stores.in_memory import InMemoryDocumentStore

document_store = InMemoryDocumentStore()

InMemoryDocumentStore 是最容易上手的 DocumentStore。它不需要任何外部依赖，是小型项目和调试的不错选择。但它扩展性不佳，不适合大型文档集合，因此不适合生产系统。要了解有关 Haystack 支持的各种外部数据库类型的更多信息，请参阅DocumentStore 集成。

获取和处理文档

您将使用 PubMed Abstracts 作为文档。Hugging Face Hub 上有许多来自 PubMed 的数据集；在本教程中，您将使用 anakin87/medrag-pubmed-chunk。

然后，您将使用简单的 for 循环从数据集中创建文档。PubMed 数据集中的每个数据点都有 4 个特征

pmid
title
content：摘要
contents：摘要 + 标题

对于搜索，您将使用 contents 特征。其他特征将作为元数据存储，您将使用它们来获得搜索结果的漂亮打印或用于元数据过滤。

from datasets import load_dataset
from haystack import Document

dataset = load_dataset("anakin87/medrag-pubmed-chunk", split="train")

docs = []
for doc in dataset:
    docs.append(
        Document(content=doc["contents"], meta={"title": doc["title"], "abstract": doc["content"], "pmid": doc["id"]})
    )

使用管道索引文档

创建一个管道，使用其嵌入来存储文档到文档存储中。对于此管道，您需要一个 DocumentSplitter 将文档分割成 512 个单词的块，一个 SentenceTransformersDocumentEmbedder 来创建用于密集检索的文档嵌入，以及一个 DocumentWriter 将文档写入文档存储。

作为嵌入模型，您将使用 Hugging Face 上的 BAAI/bge-small-en-v1.5。您可以随时测试 Hugging Face 上的其他模型，或使用其他 Embedder 来切换模型提供商。

如果此步骤对您来说太长，请将嵌入模型替换为较小的模型，例如 sentence-transformers/all-MiniLM-L6-v2 或 sentence-transformers/all-mpnet-base-v2。请确保根据您模型的令牌限制更新 split_length。

from haystack.components.writers import DocumentWriter
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
from haystack.components.preprocessors.document_splitter import DocumentSplitter
from haystack import Pipeline
from haystack.utils import ComponentDevice

document_splitter = DocumentSplitter(split_by="word", split_length=512, split_overlap=32)
document_embedder = SentenceTransformersDocumentEmbedder(
    model="BAAI/bge-small-en-v1.5", device=ComponentDevice.from_str("cuda:0")
)
document_writer = DocumentWriter(document_store)

indexing_pipeline = Pipeline()
indexing_pipeline.add_component("document_splitter", document_splitter)
indexing_pipeline.add_component("document_embedder", document_embedder)
indexing_pipeline.add_component("document_writer", document_writer)

indexing_pipeline.connect("document_splitter", "document_embedder")
indexing_pipeline.connect("document_embedder", "document_writer")

indexing_pipeline.run({"document_splitter": {"documents": docs}})

文档及其嵌入已存储在 InMemoryDocumentStore 中，现在是时候创建混合检索管道了 ✅

创建用于混合检索的管道

混合检索是指结合多种检索方法以提高整体性能。在搜索系统的上下文中，混合检索管道执行传统的基于关键词的搜索和密集向量搜索，然后使用交叉编码器模型对结果进行排序。这种组合允许搜索系统利用不同方法的优点，提供更准确和多样化的结果。

以下是混合检索管道所需的步骤

1) 初始化检索器和嵌入器

初始化一个 InMemoryEmbeddingRetriever 和 InMemoryBM25Retriever 来执行密集检索和基于关键词的检索。对于密集检索，您还需要一个 SentenceTransformersTextEmbedder，它使用与索引管道中相同的嵌入模型 BAAI/bge-small-en-v1.5 来计算搜索查询的嵌入。

from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
from haystack.components.embedders import SentenceTransformersTextEmbedder

text_embedder = SentenceTransformersTextEmbedder(
    model="BAAI/bge-small-en-v1.5", device=ComponentDevice.from_str("cuda:0")
)
embedding_retriever = InMemoryEmbeddingRetriever(document_store)
bm25_retriever = InMemoryBM25Retriever(document_store)

2) 连接检索结果

Haystack 在 DocumentJoiner 中提供了多种连接方法，可用于不同的用例，例如 merge 和 reciprocal_rank_fusion。在此示例中，您将使用默认的 concatenate 模式来连接来自两个检索器的文档，因为 Ranker 将是排序文档相关性的主要组件。

from haystack.components.joiners import DocumentJoiner

document_joiner = DocumentJoiner()

3) 对结果进行排序

使用 TransformersSimilarityRanker，它使用交叉编码器模型对给定搜索查询的所有检索文档的相关性进行评分。在此示例中，您将使用 BAAI/bge-reranker-base 模型对检索到的文档进行排序，但您可以将此模型替换为 Hugging Face 上的其他交叉编码器模型。

from haystack.components.rankers import TransformersSimilarityRanker

ranker = TransformersSimilarityRanker(model="BAAI/bge-reranker-base")

4) 创建混合检索管道

将所有初始化的组件添加到您的管道并连接它们。

from haystack import Pipeline

hybrid_retrieval = Pipeline()
hybrid_retrieval.add_component("text_embedder", text_embedder)
hybrid_retrieval.add_component("embedding_retriever", embedding_retriever)
hybrid_retrieval.add_component("bm25_retriever", bm25_retriever)
hybrid_retrieval.add_component("document_joiner", document_joiner)
hybrid_retrieval.add_component("ranker", ranker)

hybrid_retrieval.connect("text_embedder", "embedding_retriever")
hybrid_retrieval.connect("bm25_retriever", "document_joiner")
hybrid_retrieval.connect("embedding_retriever", "document_joiner")
hybrid_retrieval.connect("document_joiner", "ranker")

5) 可视化管道（可选）

要理解您如何构建混合检索管道，请使用管道的 draw() 方法。如果您正在 Google Colab 上运行此笔记本，生成的将保存在侧边栏的“文件”部分。

# hybrid_retrieval.draw("hybrid-retrieval.png")

测试混合检索

将查询传递给 text_embedder、bm25_retriever 和 ranker，然后运行检索管道

query = "apnea in infants"

result = hybrid_retrieval.run(
    {"text_embedder": {"text": query}, "bm25_retriever": {"query": query}, "ranker": {"query": query}}
)

漂亮打印结果

创建一个函数来打印一种“搜索页面”。

def pretty_print_results(prediction):
    for doc in prediction["documents"]:
        print(doc.meta["title"], "\t", doc.score)
        print(doc.meta["abstract"])
        print("\n", "\n")

pretty_print_results(result["ranker"])

下一步

🎉 恭喜！您已成功创建混合检索管道！

如果您想在 RAG 管道中使用这种检索方法，请参阅教程：创建您的第一个带检索增强的 QA 管道，了解后续步骤。

要及时了解最新的 Haystack 开发动态，您可以注册我们的新闻通讯或加入 Haystack Discord 社区。

感谢阅读！

创建自定义 SuperComponents

使用基于循环的自动纠错生成结构化输出