利用Cross Encoder实现高效文档重排序

最新推荐文章于 2025-09-19 08:44:38 发布

原创最新推荐文章于 2025-09-19 08:44:38 发布 · 774 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#python

部署运行你感兴趣的模型镜像

技术背景介绍

在现代的信息检索系统中，简单的基于向量的检索（Retriever）有时可能无法提供最优的结果排序。为了解决这个问题，我们可以使用Cross Encoder来对初步检索结果进行重新排序（Reranking），以提高检索结果的相关性。Cross Encoder利用语言模型的强大语义理解能力，能够更精确地比较查询与文档之间的相关性。

核心原理解析

Cross Encoder通过将查询与文档连接起来形成输入对，使用深度学习模型进行评分，从而为每个文档计算一个相关性分数。通过重新排序这些分数可以获得更符合用户意图的结果。

代码实现演示

以下示例展示了如何利用Langchain和Hugging Face的Cross Encoder模型来实现检索结果的重排序。

# 安装必要的Python包
!pip install faiss-cpu sentence_transformers

# 导入必要的库
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

# 加载并分割文档
documents = TextLoader("../../how_to/state_of_the_union.txt").load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
texts = text_splitter.split_documents(documents)

# 使用 Hugging Face 的 sentence-transformers 进行嵌入操作
embeddingsModel = HuggingFaceEmbeddings(model_name="sentence-transformers/msmarco-distilbert-dot-v5")
retriever = FAISS.from_documents(texts, embeddingsModel).as_retriever(search_kwargs={"k": 20})

# 基本检索操作
query = "What is the plan for the economy?"
docs = retriever.invoke(query)

# 使用 Cross Encoder 进行结果重排序
model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
compressor = CrossEncoderReranker(model=model, top_n=3)
compression_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)

compressed_docs = compression_retriever.invoke("What is the plan for the economy?")

# 输出结果
def pretty_print_docs(docs):
    print(f"\n{'-' * 100}\n".join([f"Document {i+1}:\n\n" + d.page_content for i, d in enumerate(docs)]))

pretty_print_docs(compressed_docs)