120 lines
4.6 KiB
Python
120 lines
4.6 KiB
Python
import logging
|
|
from typing import Any, Dict, List, Union
|
|
|
|
import numpy as np
|
|
|
|
from mem0.configs.rerankers.base import BaseRerankerConfig
|
|
from mem0.configs.rerankers.sentence_transformer import (
|
|
SentenceTransformerRerankerConfig,
|
|
)
|
|
from mem0.reranker.base import BaseReranker
|
|
|
|
try:
|
|
from sentence_transformers import CrossEncoder
|
|
SENTENCE_TRANSFORMERS_AVAILABLE = True
|
|
except ImportError:
|
|
SENTENCE_TRANSFORMERS_AVAILABLE = False
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SentenceTransformerReranker(BaseReranker):
|
|
"""Sentence Transformer based reranker implementation."""
|
|
|
|
def __init__(self, config: Union[BaseRerankerConfig, SentenceTransformerRerankerConfig, Dict]):
|
|
"""
|
|
Initialize Sentence Transformer reranker.
|
|
|
|
Args:
|
|
config: Configuration object with reranker parameters
|
|
"""
|
|
if not SENTENCE_TRANSFORMERS_AVAILABLE:
|
|
raise ImportError("sentence-transformers package is required for SentenceTransformerReranker. Install with: pip install sentence-transformers")
|
|
|
|
# Convert to SentenceTransformerRerankerConfig if needed
|
|
if isinstance(config, dict):
|
|
config = SentenceTransformerRerankerConfig(**config)
|
|
elif isinstance(config, BaseRerankerConfig) and not isinstance(config, SentenceTransformerRerankerConfig):
|
|
# Convert BaseRerankerConfig to SentenceTransformerRerankerConfig with defaults
|
|
config = SentenceTransformerRerankerConfig(
|
|
provider=getattr(config, 'provider', 'sentence_transformer'),
|
|
model=getattr(config, 'model', 'cross-encoder/ms-marco-MiniLM-L-6-v2'),
|
|
api_key=getattr(config, 'api_key', None),
|
|
top_k=getattr(config, 'top_k', None),
|
|
device=None, # Will auto-detect
|
|
batch_size=32, # Default
|
|
show_progress_bar=False, # Default
|
|
)
|
|
|
|
self.config = config
|
|
self.model = CrossEncoder(self.config.model, device=self.config.device)
|
|
|
|
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Rerank documents using sentence transformer cross-encoder.
|
|
|
|
Args:
|
|
query: The search query
|
|
documents: List of documents to rerank
|
|
top_k: Number of top documents to return
|
|
|
|
Returns:
|
|
List of reranked documents with rerank_score
|
|
"""
|
|
if not documents:
|
|
return documents
|
|
|
|
# Extract text content for reranking
|
|
doc_texts = []
|
|
for doc in documents:
|
|
if 'memory' in doc:
|
|
doc_texts.append(doc['memory'])
|
|
elif 'text' in doc:
|
|
doc_texts.append(doc['text'])
|
|
elif 'content' in doc:
|
|
doc_texts.append(doc['content'])
|
|
else:
|
|
doc_texts.append(str(doc))
|
|
|
|
try:
|
|
# Create query-document pairs
|
|
pairs = [[query, doc_text] for doc_text in doc_texts]
|
|
|
|
scores = self.model.predict(
|
|
pairs,
|
|
batch_size=self.config.batch_size,
|
|
show_progress_bar=self.config.show_progress_bar,
|
|
)
|
|
if isinstance(scores, np.ndarray):
|
|
scores = scores.tolist()
|
|
|
|
# Combine documents with scores
|
|
doc_score_pairs = list(zip(documents, scores))
|
|
|
|
# Sort by score (descending)
|
|
doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
|
|
|
|
# Apply top_k limit
|
|
final_top_k = top_k or self.config.top_k
|
|
if final_top_k:
|
|
doc_score_pairs = doc_score_pairs[:final_top_k]
|
|
|
|
# Create reranked results
|
|
reranked_docs = []
|
|
for doc, score in doc_score_pairs:
|
|
reranked_doc = doc.copy()
|
|
reranked_doc['rerank_score'] = float(score)
|
|
reranked_docs.append(reranked_doc)
|
|
|
|
return reranked_docs
|
|
|
|
except Exception as e:
|
|
# Fallback to original order if reranking fails
|
|
logger.warning("SentenceTransformer reranking failed, falling back to original order: %s", e)
|
|
fallback_docs = []
|
|
for doc in documents:
|
|
fallback_doc = doc.copy()
|
|
fallback_doc['rerank_score'] = 0.0
|
|
fallback_docs.append(fallback_doc)
|
|
final_top_k = top_k or self.config.top_k
|
|
return fallback_docs[:final_top_k] if final_top_k else fallback_docs |