高效信息检索：使用Cross Encoder进行重排序

最新推荐文章于 2024-10-07 15:49:09 发布

afTFODguAKBF

最新推荐文章于 2024-10-07 15:49:09 发布

阅读量475

点赞数 3

文章标签： python 算法开发语言

本文链接：https://blog.csdn.net/afTFODguAKBF/article/details/142408714

版权

引言

在信息检索系统中，提高检索结果的质量至关重要。Cross Encoder是一种有效的重排序机制，可与嵌入技术结合使用，以提升检索效果。这篇文章将介绍如何在向量检索系统中集成Cross Encoder，尤其是如何利用Hugging Face的Cross Encoder模型在SageMaker中实现重排序。

主要内容

向量检索的设置

我们首先使用一个简单的向量检索系统，从文本中提取信息。以下代码展示了如何初始化向量检索，并从文本中检索信息。

# 安装必要的库
#!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("path/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进行重排序

我们使用Hugging Face的Cross Encoder模型，通过SageMaker实现对检索结果的重排序。

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)

在SageMaker上部署

以下是用于在SageMaker上创建Cross Encoder端点的示例代码。

import json
import logging
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):
        with torch.inference_mode():
            inputs = self.tokenizer(
                pairs,
                padding=True,
                truncation=True,
                return_tensors="pt",
                max_length=512,
            ).to(self.device)
            
            scores = self.model(**inputs).logits.view(-1).float()
        return scores.detach().cpu().tolist()

def model_fn(model_dir: str) -> CrossEncoder:
    return CrossEncoder()

def transform_fn(cross_encoder, input_data, content_type, accept):
    payload = json.loads(input_data)
    model_output = cross_encoder(**payload)
    output = {SCORES: model_output}
    return encoder.encode(output, accept)