fix(reranker): score HuggingFace reranker with sigmoid, not min-max (#5715)
This commit is contained in:
@@ -59,6 +59,26 @@ class HuggingFaceReranker(BaseReranker):
|
||||
self.model.to(self.device)
|
||||
self.model.eval()
|
||||
|
||||
@staticmethod
|
||||
def _normalize_scores(scores: List[float]) -> List[float]:
|
||||
"""Map raw cross-encoder logits into the [0, 1] range via a sigmoid.
|
||||
|
||||
Cross-encoder rerankers (e.g. ``BAAI/bge-reranker-*``) emit unbounded
|
||||
logits; the documented way to obtain an interpretable [0, 1] relevance
|
||||
score is a per-document sigmoid (``1 / (1 + e^-x)``), which preserves the
|
||||
ranking order.
|
||||
|
||||
This replaces the previous min-max scaling, which produced *set-relative*
|
||||
scores: the lowest-ranked document was always forced to 0.0, and a single
|
||||
document (or any set of tied scores) collapsed to 0.0 — wrongly reporting
|
||||
a result as completely irrelevant. Sigmoid scores each document on its own
|
||||
merit, so those cases are handled naturally.
|
||||
"""
|
||||
if not scores:
|
||||
return []
|
||||
arr = np.asarray(scores, dtype=float)
|
||||
return (1.0 / (1.0 + np.exp(-arr))).tolist()
|
||||
|
||||
def rerank(self, query: str, documents: List[Dict[str, Any]], top_k: int = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Rerank documents using HuggingFace cross-encoder model.
|
||||
@@ -118,9 +138,7 @@ class HuggingFaceReranker(BaseReranker):
|
||||
|
||||
# Normalize scores if requested
|
||||
if self.config.normalize:
|
||||
scores = np.array(scores)
|
||||
scores = (scores - scores.min()) / (scores.max() - scores.min() + 1e-8)
|
||||
scores = scores.tolist()
|
||||
scores = self._normalize_scores(scores)
|
||||
|
||||
# Combine documents with scores
|
||||
doc_score_pairs = list(zip(documents, scores))
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Unit tests for HuggingFaceReranker score normalization.
|
||||
|
||||
These exercise the pure ``_normalize_scores`` helper directly, so they do not
|
||||
require ``transformers`` / ``torch`` to be installed.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
from mem0.reranker.huggingface_reranker import HuggingFaceReranker
|
||||
|
||||
|
||||
def _sigmoid(x):
|
||||
return 1.0 / (1.0 + math.exp(-x))
|
||||
|
||||
|
||||
class TestHuggingFaceNormalizeScores:
|
||||
def test_logits_mapped_via_sigmoid(self):
|
||||
scores = HuggingFaceReranker._normalize_scores([2.0, 8.0, 5.0])
|
||||
assert scores == pytest.approx([_sigmoid(2.0), _sigmoid(8.0), _sigmoid(5.0)])
|
||||
|
||||
def test_output_bounded_between_zero_and_one(self):
|
||||
for s in HuggingFaceReranker._normalize_scores([-12.0, -1.0, 0.0, 3.0, 15.0]):
|
||||
assert 0.0 <= s <= 1.0
|
||||
|
||||
def test_sigmoid_preserves_ranking_order(self):
|
||||
raw = [1.0, -4.0, 9.0, 2.5]
|
||||
normalized = HuggingFaceReranker._normalize_scores(raw)
|
||||
# argsort of raw and normalized must match — sigmoid is monotonic.
|
||||
assert sorted(range(len(raw)), key=lambda i: raw[i]) == sorted(
|
||||
range(len(normalized)), key=lambda i: normalized[i]
|
||||
)
|
||||
|
||||
def test_single_score_not_collapsed_to_zero(self):
|
||||
# Regression: a lone document used to normalize to ~0.0 under min-max.
|
||||
# A positive logit must now yield a clearly-relevant score (> 0.5).
|
||||
(score,) = HuggingFaceReranker._normalize_scores([4.2])
|
||||
assert score == pytest.approx(_sigmoid(4.2))
|
||||
assert score > 0.5
|
||||
|
||||
def test_tied_scores_not_collapsed_to_zero(self):
|
||||
# Regression: tied candidates all collapsed to ~0.0 under min-max.
|
||||
scores = HuggingFaceReranker._normalize_scores([3.0, 3.0, 3.0])
|
||||
assert scores == pytest.approx([_sigmoid(3.0)] * 3)
|
||||
|
||||
def test_zero_logit_maps_to_half(self):
|
||||
assert HuggingFaceReranker._normalize_scores([0.0]) == pytest.approx([0.5])
|
||||
|
||||
def test_empty_scores(self):
|
||||
assert HuggingFaceReranker._normalize_scores([]) == []
|
||||
Reference in New Issue
Block a user