diff --git a/hindsight-api-slim/hindsight_api/engine/search/reranking.py b/hindsight-api-slim/hindsight_api/engine/search/reranking.py index 4888400b69..4698626638 100644 --- a/hindsight-api-slim/hindsight_api/engine/search/reranking.py +++ b/hindsight-api-slim/hindsight_api/engine/search/reranking.py @@ -212,16 +212,47 @@ async def rerank(self, query: str, candidates: list[MergedCandidate]) -> list[Sc # Get cross-encoder scores scores = await self.cross_encoder.predict(pairs) - # Normalize scores using sigmoid to [0, 1] range - # Cross-encoder returns logits which can be negative - import math - + # Normalize scores to [0, 1] range. + # Local models return logits (any real number) — sigmoid is appropriate. + # External API rerankers (SiliconFlow, Cohere, etc.) return pre-normalized + # relevance_score in [0, 1] with very small absolute values. Applying + # sigmoid to these compresses everything to ~0.5, destroying the ranking + # signal and making recency the sole sorting factor. We detect the score range + # and choose the appropriate normalization. import numpy as np - def sigmoid(x): + def _sigmoid(x: float) -> float: return 1 / (1 + np.exp(-x)) - normalized_scores = [sigmoid(score) for score in scores] + def _rank_normalize_with_ties(score_list: list[float]) -> list[float]: + """Rank-based normalization that assigns equal ranks to equal scores.""" + n = len(score_list) + if n <= 1: + return [1.0] * n + indexed = sorted(enumerate(score_list), key=lambda x: x[1], reverse=True) + result = [0.0] * n + i = 0 + while i < n: + j = i + while j < n and indexed[j][1] == indexed[i][1]: + j += 1 + # Average rank for tied scores + avg_rank = (i + j - 1) / 2.0 + norm = 1.0 - (avg_rank / (n - 1)) + for k in range(i, j): + result[indexed[k][0]] = norm + i = j + return result + + if scores and min(scores) >= 0.0 and max(scores) <= 1.0: + # Scores are already in [0, 1] (e.g. SiliconFlow, Cohere relevance_score). + # Use rank-based normalization to preserve relative ordering without + # depending on absolute score magnitudes. + normalized_scores = _rank_normalize_with_ties(scores) + else: + # Scores are logits (e.g. local sentence-transformers models). + # Sigmoid maps (-inf, +inf) to (0, 1). + normalized_scores = [_sigmoid(score) for score in scores] # Create ScoredResult objects with cross-encoder scores scored_results = [] diff --git a/hindsight-api-slim/tests/test_reranker_score_normalization.py b/hindsight-api-slim/tests/test_reranker_score_normalization.py new file mode 100644 index 0000000000..fad4aeac64 --- /dev/null +++ b/hindsight-api-slim/tests/test_reranker_score_normalization.py @@ -0,0 +1,163 @@ +""" +Unit tests for CrossEncoderReranker.rerank() score normalization logic. + +Covers: +1. Rank-based normalization when scores are already in [0, 1]. +2. Tied scores receiving identical normalized values. +3. Sigmoid normalization when scores are logits outside [0, 1]. +4. Empty candidates returning an empty list without calling predict(). +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest + +from hindsight_api.engine.search.reranking import CrossEncoderReranker +from hindsight_api.engine.search.types import MergedCandidate, RetrievalResult + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_candidates(n: int) -> list[MergedCandidate]: + """Create *n* minimal MergedCandidate objects.""" + candidates = [] + for i in range(n): + retrieval = RetrievalResult( + id=str(uuid4()), + text=f"Document {i}", + fact_type="world", + occurred_start=None, + occurred_end=None, + ) + candidates.append( + MergedCandidate(retrieval=retrieval, rrf_score=1.0 / (i + 1)) + ) + return candidates + + +def _make_cross_encoder(predict_return: list[float]): + """Return a fake cross-encoder whose `predict` is an AsyncMock.""" + ce = AsyncMock() + ce.predict = AsyncMock(return_value=predict_return) + ce.provider_name = "local" + ce.initialize = AsyncMock() + return ce + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_rank_normalization_for_0_1_scores(): + """Scores already in [0, 1] should use rank-based normalization.""" + # Scores in [0, 1] but NOT logit-shaped — rank normalization preserves order + raw_scores = [0.1, 0.5, 0.9] + ce = _make_cross_encoder(raw_scores) + reranker = CrossEncoderReranker(cross_encoder=ce) + reranker._initialized = True + + candidates = _make_candidates(3) + results = await reranker.rerank("test query", candidates) + + assert len(results) == 3 + # Top-ranked (0.9) gets normalized 1.0 + assert results[0].cross_encoder_score == pytest.approx(0.9) + assert results[0].cross_encoder_score_normalized == pytest.approx(1.0) + # Bottom-ranked (0.1) gets normalized 0.0 + assert results[-1].cross_encoder_score == pytest.approx(0.1) + assert results[-1].cross_encoder_score_normalized == pytest.approx(0.0) + # Middle gets 0.5 + assert results[1].cross_encoder_score_normalized == pytest.approx(0.5) + # predict was called exactly once + ce.predict.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_tied_scores_get_same_normalized_value(): + """When raw scores contain ties, tied entries must receive the same normalized score.""" + # Two scores tied at 0.7, one distinct at 0.3 + raw_scores = [0.7, 0.3, 0.7] + ce = _make_cross_encoder(raw_scores) + reranker = CrossEncoderReranker(cross_encoder=ce) + reranker._initialized = True + + candidates = _make_candidates(3) + results = await reranker.rerank("test query", candidates) + + # The two 0.7 entries should share the same normalized value. + # With rank-based normalization over 3 items (indices 0,1,2 sorted desc): + # rank 0 & 1 tied → avg_rank = 0.5 → norm = 1 - 0.5/2 = 0.75 + # rank 2 → norm = 1 - 2/2 = 0.0 + # After sorting by normalized score descending, the two 0.7 entries are + # at indices 0 and 1; the 0.3 entry is at index 2. + scores_by_raw = {r.cross_encoder_score: r.cross_encoder_score_normalized for r in results} + tied_norm = scores_by_raw[0.7] + assert tied_norm == pytest.approx(0.75) + assert scores_by_raw[0.3] == pytest.approx(0.0) + # Both 0.7 results have identical normalized scores + assert results[0].cross_encoder_score_normalized == pytest.approx( + results[1].cross_encoder_score_normalized + ) + + +@pytest.mark.asyncio +async def test_sigmoid_normalization_for_logits(): + """When scores are outside [0, 1] (logits), sigmoid normalization is used.""" + raw_scores = [2.0, -1.0, 0.0] + ce = _make_cross_encoder(raw_scores) + reranker = CrossEncoderReranker(cross_encoder=ce) + reranker._initialized = True + + candidates = _make_candidates(3) + results = await reranker.rerank("test query", candidates) + + assert len(results) == 3 + import math + + # Results are sorted by weight descending + expected_sigmoid = [1 / (1 + math.exp(-s)) for s in raw_scores] + expected_sorted = sorted(expected_sigmoid, reverse=True) + + for result, expected in zip(results, expected_sorted): + assert result.cross_encoder_score_normalized == pytest.approx(expected, rel=1e-6) + + # Verify the highest logit (2.0) maps to the highest normalized score + assert results[0].cross_encoder_score == pytest.approx(2.0) + assert results[0].cross_encoder_score_normalized > 0.5 + + +@pytest.mark.asyncio +async def test_empty_candidates_returns_empty_without_predict(): + """When candidates are empty, rerank must return [] without calling predict().""" + ce = _make_cross_encoder([]) + reranker = CrossEncoderReranker(cross_encoder=ce) + reranker._initialized = True + + results = await reranker.rerank("test query", []) + + assert results == [] + ce.predict.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_mixed_boundary_scores_use_rank(): + """Boundary values 0.0 and 1.0 (still in [0,1]) should trigger rank normalization.""" + raw_scores = [0.0, 1.0, 0.5] + ce = _make_cross_encoder(raw_scores) + reranker = CrossEncoderReranker(cross_encoder=ce) + reranker._initialized = True + + candidates = _make_candidates(3) + results = await reranker.rerank("test query", candidates) + + # Rank-based: 1.0 → 1.0, 0.5 → 0.5, 0.0 → 0.0 + by_score = {r.cross_encoder_score: r.cross_encoder_score_normalized for r in results} + assert by_score[1.0] == pytest.approx(1.0) + assert by_score[0.5] == pytest.approx(0.5) + assert by_score[0.0] == pytest.approx(0.0)