JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix incorrect use of SentenceTransformer for cross-encoder reranker models (#4806)

Browse Source
This commit is contained in:
Chinnu Abey
2026-04-13 20:14:22 +05:30
committed by GitHub
parent cf530c4bec
commit ac9cdd4840
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/reranker/sentence_transformer_reranker.py
+7 -4
View File
@@ -6,7 +6,7 @@ from mem0.configs.rerankers.base import BaseRerankerConfig
from mem0.configs.rerankers.sentence_transformer import SentenceTransformerRerankerConfig from mem0.configs.rerankers.sentence_transformer import SentenceTransformerRerankerConfig
try: try:
from sentence_transformers import SentenceTransformer from sentence_transformers import CrossEncoder
SENTENCE_TRANSFORMERS_AVAILABLE = True SENTENCE_TRANSFORMERS_AVAILABLE = True
except ImportError: except ImportError:
SENTENCE_TRANSFORMERS_AVAILABLE = False SENTENCE_TRANSFORMERS_AVAILABLE = False
@@ -41,7 +41,7 @@ class SentenceTransformerReranker(BaseReranker):
) )
self.config = config self.config = config
self.model = SentenceTransformer(self.config.model, device=self.config.device) 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]]: def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
""" """
@@ -74,8 +74,11 @@ class SentenceTransformerReranker(BaseReranker):
# Create query-document pairs # Create query-document pairs
pairs = [[query, doc_text] for doc_text in doc_texts] pairs = [[query, doc_text] for doc_text in doc_texts]
# Get similarity scores scores = self.model.predict(
scores = self.model.predict(pairs) pairs,
batch_size=self.config.batch_size,
show_progress_bar=self.config.show_progress_bar,
)
if isinstance(scores, np.ndarray): if isinstance(scores, np.ndarray):
scores = scores.tolist() scores = scores.tolist()