040eb33ea6
* improve ci and tests * add more tests * fixes * fix tests * fix more * fix * fix * fix * fix * fix for real * tests and doc * fix cp * fix link pg0 * fix pg0 * fix pg0 * fix pg0 * even better * more * fix
102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
"""
|
|
Cross-encoder abstraction for reranking.
|
|
|
|
Provides an interface for reranking with different backends.
|
|
"""
|
|
from abc import ABC, abstractmethod
|
|
from typing import List, Tuple
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CrossEncoderModel(ABC):
|
|
"""
|
|
Abstract base class for cross-encoder reranking.
|
|
|
|
Cross-encoders take query-document pairs and return relevance scores.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def load(self) -> None:
|
|
"""
|
|
Load the cross-encoder model.
|
|
|
|
This should be called during initialization to load the model
|
|
and avoid cold start latency on first predict() call.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
|
|
"""
|
|
Score query-document pairs for relevance.
|
|
|
|
Args:
|
|
pairs: List of (query, document) tuples to score
|
|
|
|
Returns:
|
|
List of relevance scores (higher = more relevant)
|
|
"""
|
|
pass
|
|
|
|
|
|
class SentenceTransformersCrossEncoder(CrossEncoderModel):
|
|
"""
|
|
Cross-encoder implementation using SentenceTransformers.
|
|
|
|
Call load() during initialization to load the model and avoid cold starts.
|
|
|
|
Default model is cross-encoder/ms-marco-MiniLM-L-6-v2:
|
|
- Fast inference (~80ms for 100 pairs on CPU)
|
|
- Small model (80MB)
|
|
- Trained for passage re-ranking
|
|
"""
|
|
|
|
def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
|
|
"""
|
|
Initialize SentenceTransformers cross-encoder.
|
|
|
|
Args:
|
|
model_name: Name of the CrossEncoder model to use.
|
|
Default: cross-encoder/ms-marco-MiniLM-L-6-v2
|
|
"""
|
|
self.model_name = model_name
|
|
self._model = None
|
|
|
|
def load(self) -> None:
|
|
"""Load the cross-encoder model."""
|
|
if self._model is not None:
|
|
return
|
|
|
|
try:
|
|
from sentence_transformers import CrossEncoder
|
|
except ImportError:
|
|
raise ImportError(
|
|
"sentence-transformers is required for SentenceTransformersCrossEncoder. "
|
|
"Install it with: pip install sentence-transformers"
|
|
)
|
|
|
|
logger.info(f"Loading cross-encoder model: {self.model_name}...")
|
|
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
|
|
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
|
|
self._model = CrossEncoder(
|
|
self.model_name,
|
|
model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
|
|
)
|
|
logger.info("Cross-encoder model loaded")
|
|
|
|
def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
|
|
"""
|
|
Score query-document pairs for relevance.
|
|
|
|
Args:
|
|
pairs: List of (query, document) tuples to score
|
|
|
|
Returns:
|
|
List of relevance scores (raw logits from the model)
|
|
"""
|
|
if self._model is None:
|
|
self.load()
|
|
scores = self._model.predict(pairs, show_progress_bar=False)
|
|
return scores.tolist() if hasattr(scores, 'tolist') else list(scores)
|