80 lines
2.4 KiB
Python
80 lines
2.4 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 CrossEncoderReranker(ABC):
|
|
"""
|
|
Abstract base class for cross-encoder reranking.
|
|
|
|
Cross-encoders take query-document pairs and return relevance scores.
|
|
"""
|
|
|
|
@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(CrossEncoderReranker):
|
|
"""
|
|
Cross-encoder implementation using SentenceTransformers.
|
|
|
|
Uses lazy import so sentence-transformers is not required if another
|
|
reranking backend is used.
|
|
|
|
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 and load model.
|
|
|
|
Args:
|
|
model_name: Name of the CrossEncoder model to use.
|
|
Default: cross-encoder/ms-marco-MiniLM-L-6-v2
|
|
"""
|
|
self.model_name = model_name
|
|
|
|
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}...")
|
|
self._model = CrossEncoder(self.model_name)
|
|
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)
|
|
"""
|
|
scores = self._model.predict(pairs)
|
|
return scores.tolist() if hasattr(scores, 'tolist') else list(scores)
|