提升搜索引擎性能：使用Cross Encoder Reranker优化结果排序

本文链接：https://blog.youkuaiyun.com/tt_jishu/article/details/143421060

# 介绍

在现代信息检索系统中，检索器（Retriever）和重排序器（Reranker）通常被结合使用，以提高搜索结果的质量。Cross Encoder Reranker是一种重排序技术，通过使用模型对检索的初步文档进行重新排序，以更好地匹配查询意图。在本文中，我们将详细介绍如何使用Hugging Face的Cross Encoder模型实现Cross Encoder Reranker，并展示如何在Amazon SageMaker上部署这些模型。

# 主要内容

## 1. 向量存储检索器的设置

我们首先设置一个简单的向量存储检索器，利用2023年美国国情咨文的文本数据。该检索器将会检索出20篇相关文档。

```python
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)
pretty_print_docs(docs)

2. 用CrossEncoderReranker进行重排序

我们将使用ContextualCompressionRetriever和CrossEncoderReranker对检索结果进行重排序，从而提高结果的质量。

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

# 配置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?")
pretty_print_docs(compressed_docs)

3. 在SageMaker中部署Cross Encoder模型

通过在Amazon SageMaker上部署Cross Encoder模型，您可以在生产环境中使用它们以获得高效的推理能力。以下是一个简单的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 torch.device("cpu")
        logging.info(f"Using device: {self.device}")
        model_name = "BAAI/bge-reranker-base"
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model = 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 = inputs.to(self.device)
            scores = (
                self.model(**inputs, return_dict=True).logits.view(-1).float()
            )

        return scores.detach().cpu().tolist()

def model_fn(model_dir: str) -> CrossEncoder:
    try:
        return CrossEncoder()
    except Exception:
        logging.exception(f"Failed to load model from: {model_dir}")
        raise

def transform_fn(cross_encoder: CrossEncoder, input_data: bytes, content_type: str, accept: str) -> bytes:
    payload = json.loads(input_data)
    model_output = cross_encoder(**payload)
    output = {"scores": model_output}
    return encoder.encode(output, accept)