【langchain学习】BM25Retriever和FaissRetriever组合实现EnsembleRetriever混合检索器的实践

展示如何使用 LangChain 的 EnsembleRetriever 组合 BM25 和 FAISS 两种检索方法，从而在检索过程中结合关键词匹配和语义相似性搜索的优势。通过这种组合，我们能够在查询时获得更全面的结果。

1. 导入必要的库和模块

首先，我们需要导入所需的库和模块。本文使用的工具包括 LangChain 提供的 EnsembleRetriever、BM25Retriever、FAISS、以及 HuggingFace 的嵌入模型。

from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from config import paraphrase_multilingual_MiniLM_L12_v2

2. 准备文档数据

准备了两组文档数据，这些文档分别用于 BM25 和 FAISS 检索器。在实际使用中，这些文档可以是任何你想要检索的文本数据。

# 定义第一组文档，这些文档将用于 BM25 检索器
doc_list_1 = [
    "这是一个测试句子",
    "温格高赢得了2023环法冠军",
    "波士顿马拉松是历史悠久的一项比赛",
    "何杰即将出战巴黎奥运会的马拉松项目",
    "珍宝将不再赞助温格高所在的车队",
]

# 定义第二组文档，这些文档将用于 FAISS 检索器
doc_list_2 = [
    "波加查擅长陡坡进攻，而温格高则更擅长长坡",
    "温格高的最大摄氧量居然有97！",
    "北京奥运会在2008年8月8日开幕",
    "基普乔格是东京马拉松的金牌得主",
]

3. 初始化 BM25 检索器

BM25 是一种基于词频和逆文档频率（TF-IDF）的传统检索算法，非常适合关键词匹配。我们使用 BM25Retriever.from_texts 方法来创建 BM25 检索器，并为其设置元数据（source: 1），以便区分文档来源。我们还设置了返回的文档数量 k 为 2。

bm25_retriever = BM25Retriever.from_texts(
    doc_list_1, metadatas=[{"source": 1}] * len(doc_list_1)
)
bm25_retriever.k = 2  # 设置 BM25 检索器返回的文档数量

4. 初始化 FAISS 检索器

FAISS 是一种用于高效向量相似性搜索的工具。我们使用 HuggingFace 的 paraphrase_multilingual_MiniLM_L12_v2 嵌入模型将文档转化为向量，然后将这些向量存储在 FAISS 中。随后，我们使用 FAISS.from_texts 方法创建 FAISS 检索器，并设置返回的文档数量 k 为 2。

embeddings = HuggingFaceEmbeddings(model_name=paraphrase_multilingual_MiniLM_L12_v2)
faiss_vectorstore = FAISS.from_texts(
    doc_list_2, embeddings, metadatas=[{"source": 2}] * len(doc_list_2)
)
faiss_retriever = faiss_vectorstore.as_retriever(search_kwargs={"k": 2})

5. 创建组合检索器

为了结合 BM25 和 FAISS 的优势，我们使用 EnsembleRetriever 将这两种检索器组合在一起，并为它们设置相等的权重（0.5 和 0.5）。这种组合策略可以在关键词匹配和语义相似性之间取得平衡。

ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, faiss_retriever], weights=[0.5, 0.5]
)

6. 执行查询并输出结果

我们使用组合检索器来查询关键词“温格高”，并输出检索结果。EnsembleRetriever 会结合 BM25 和 FAISS 的结果，返回与查询最相关的文档。随后，我们提取并打印每个文档的 page_content。

docs = ensemble_retriever.invoke("温格高")
print(docs)

page_contents = [doc.page_content for doc in docs]
print(page_contents)

7. 检索结果分析

运行上述代码后，我们得到了如下结果：

[Document(page_content='珍宝将不再赞助温格高所在的车队', metadata={'source': 1}),
 Document(page_content='温格高的最大摄氧量居然有97！', metadata={'source': 2}),
 Document(page_content='温格高赢得了2023环法冠军', metadata={'source': 1}),
 Document(page_content='波加查擅长陡坡进攻，而温格高则更擅长长坡', metadata={'source': 2})]

['珍宝将不再赞助温格高所在的车队', '温格高的最大摄氧量居然有97！', '温格高赢得了2023环法冠军', '波加查擅长陡坡进攻，而温格高则更擅长长坡']

这些结果显示，组合检索器从两个检索器中分别返回了两个文档，并将它们合并。最终结果中既包含了 BM25 检索器基于词频的匹配结果，也包含了 FAISS 检索器基于语义相似性的结果。例如，“珍宝将不再赞助温格高所在的车队”和“温格高的最大摄氧量居然有97！”这两条结果分别来自 BM25 和 FAISS，它们都与查询词“温格高”高度相关。