141 lines
4.1 KiB
Python
141 lines
4.1 KiB
Python
"""
|
||
Cross-Encoder 重排序器
|
||
|
||
使用 sentence-transformers 加载交叉编码器模型，对检索结果进行精排。
|
||
"""
|
||
|
||
import os
|
||
from typing import List, Dict, Any, Optional
|
||
from langchain.retrievers.document_compressors import CrossEncoderReranker
|
||
from langchain_core.documents import Document
|
||
from sentence_transformers import CrossEncoder
|
||
|
||
|
||
class CrossEncoderReranker:
|
||
"""
|
||
Cross-Encoder 重排序器包装类
|
||
|
||
支持 BAAI/bge-reranker-base 等模型。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model_name: str = "BAAI/bge-reranker-base",
|
||
top_n: int = 5,
|
||
device: Optional[str] = None,
|
||
cache_folder: Optional[str] = None,
|
||
):
|
||
"""
|
||
初始化重排序器
|
||
|
||
Args:
|
||
model_name: 模型名称或路径
|
||
top_n: 返回的顶部文档数量
|
||
device: 设备（cpu/cuda），如果为 None 则自动选择
|
||
cache_folder: 模型缓存目录
|
||
"""
|
||
self.model_name = model_name
|
||
self.top_n = top_n
|
||
self.device = device
|
||
self.cache_folder = cache_folder or os.path.join(
|
||
os.path.expanduser("~"), ".cache", "sentence_transformers"
|
||
)
|
||
|
||
# 延迟加载模型
|
||
self._model = None
|
||
self._langchain_reranker = None
|
||
|
||
def _load_model(self):
|
||
"""加载交叉编码器模型"""
|
||
if self._model is None:
|
||
try:
|
||
self._model = CrossEncoder(
|
||
self.model_name,
|
||
device=self.device,
|
||
cache_folder=self.cache_folder,
|
||
)
|
||
except Exception as e:
|
||
# 如果指定模型加载失败，尝试备用模型
|
||
print(f"加载模型 {self.model_name} 失败: {e}")
|
||
print("尝试加载备用模型 BAAI/bge-reranker-v2-m3...")
|
||
self._model = CrossEncoder(
|
||
"BAAI/bge-reranker-v2-m3",
|
||
device=self.device,
|
||
cache_folder=self.cache_folder,
|
||
)
|
||
|
||
def _create_langchain_reranker(self):
|
||
"""创建 LangChain 重排序器"""
|
||
if self._langchain_reranker is None:
|
||
self._load_model()
|
||
self._langchain_reranker = CrossEncoderReranker(
|
||
model=self._model,
|
||
top_n=self.top_n,
|
||
)
|
||
|
||
def rerank(
|
||
self,
|
||
query: str,
|
||
documents: List[Document],
|
||
) -> List[Document]:
|
||
"""
|
||
对文档进行重排序
|
||
|
||
Args:
|
||
query: 查询文本
|
||
documents: 待排序文档列表
|
||
|
||
Returns:
|
||
重排序后的文档列表
|
||
"""
|
||
self._create_langchain_reranker()
|
||
return self._langchain_reranker.compress_documents(
|
||
documents=documents,
|
||
query=query,
|
||
)
|
||
|
||
def create_contextual_compression_retriever(
|
||
self,
|
||
base_retriever: Any,
|
||
) -> Any:
|
||
"""
|
||
创建上下文压缩检索器
|
||
|
||
Args:
|
||
base_retriever: 基础检索器
|
||
|
||
Returns:
|
||
上下文压缩检索器
|
||
"""
|
||
from langchain.retrievers import ContextualCompressionRetriever
|
||
|
||
self._create_langchain_reranker()
|
||
|
||
compression_retriever = ContextualCompressionRetriever(
|
||
base_compressor=self._langchain_reranker,
|
||
base_retriever=base_retriever,
|
||
)
|
||
|
||
return compression_retriever
|
||
|
||
@classmethod
|
||
def create_from_config(
|
||
cls,
|
||
config: Optional[Dict[str, Any]] = None,
|
||
) -> "CrossEncoderReranker":
|
||
"""
|
||
从配置创建重排序器
|
||
|
||
Args:
|
||
config: 配置字典，包含 model_name, top_n, device 等
|
||
|
||
Returns:
|
||
CrossEncoderReranker 实例
|
||
"""
|
||
config = config or {}
|
||
return cls(
|
||
model_name=config.get("model_name", "BAAI/bge-reranker-base"),
|
||
top_n=config.get("top_n", 5),
|
||
device=config.get("device", None),
|
||
cache_folder=config.get("cache_folder", None),
|
||
) |