使用Cross Encoder进行重排序:从Retriever到SageMaker端点

老铁们,今天我们来聊聊如何使用Cross Encoder进行重排序(reranking)操作。我们会从一个基本的检索器出发,然后用Hugging Face提供的Cross Encoder模型来实现重排序。说白了,这个过程就是通过引入Cross Encoder来提高检索的准确性。在这里,我会用到Hugging Face的模型,比如BAAI/bge-reranker-base,并在SageMaker中部署。

技术背景介绍

在信息检索系统中,重排序是提高检索精度的重要技术手段。Cross Encoder可以与嵌入结合使用,作为一种强大的重排序机制。Hugging Face的Cross-Encoders文档对此有详细说明。

原理深度解析

Cross Encoder的工作原理是对文本对进行编码,然后通过模型计算出相关性得分。相比于单独使用嵌入,Cross Encoder考虑了句子间的关系,因此能提供更精准的排序结果。

实战代码演示

设置基础向量存储检索器

首先,我们需要初始化一个简单的向量存储检索器,并存储文档。例如,这里我们用到了2023年的国情咨文作为示例数据。

#!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)
pretty_print_docs(docs)
使用CrossEncoder进行重排序

接下来,我们在基础检索器上包装一个ContextualCompressionRetriever,并利用CrossEncoderReranker进行重排序。

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

优化建议分享

我先前踩过这个坑,使用Cross Encoder进行重排序虽然效果显著,但计算量较大。建议在生产环境中使用代理服务提高稳定性,比如把模型部署到AWS SageMaker上。

补充说明和总结

在AWS SageMaker上部署Hugging Face的Cross Encoder模型可以利用其强大的计算能力,具体可以参考以下示例inference.py来创建SageMaker端点。代码下载模型,并在端点上实时运行推理。

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

PAIRS = "pairs"
SCORES = "scores"

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)

今天的技术分享就到这里,希望对大家有帮助。开发过程中遇到问题也可以在评论区交流~

—END—

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值