Skip to content

Commit dcf5588

Browse files
xuli500177root
andauthored
fix(reranker): detect pre-normalized scores and use rank-based normalization (#1512)
* fix(reranker): detect pre-normalized scores and use rank-based normalization 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. This fix detects the score range: - If all scores are in [0, 1]: use rank-based normalization with tie handling (equal scores get equal ranks) - Otherwise (logits): use sigmoid as before This preserves the correct behavior for local models (logits) while fixing ranking quality for external API rerankers. * test(reranker): add unit tests for score normalization logic - Rank-based normalization for [0,1] scores - Tied scores receive identical normalized values - Sigmoid normalization for logit scores - Empty candidates returns [] without calling predict() - Fix typo: "sole排序 factor" -> "sole sorting factor" --------- Co-authored-by: root <root@LAPTOP-4VPLUOA6.localdomain>
1 parent 3d6c2ba commit dcf5588

2 files changed

Lines changed: 200 additions & 6 deletions

File tree

hindsight-api-slim/hindsight_api/engine/search/reranking.py

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -212,16 +212,47 @@ async def rerank(self, query: str, candidates: list[MergedCandidate]) -> list[Sc
212212
# Get cross-encoder scores
213213
scores = await self.cross_encoder.predict(pairs)
214214

215-
# Normalize scores using sigmoid to [0, 1] range
216-
# Cross-encoder returns logits which can be negative
217-
import math
218-
215+
# Normalize scores to [0, 1] range.
216+
# Local models return logits (any real number) — sigmoid is appropriate.
217+
# External API rerankers (SiliconFlow, Cohere, etc.) return pre-normalized
218+
# relevance_score in [0, 1] with very small absolute values. Applying
219+
# sigmoid to these compresses everything to ~0.5, destroying the ranking
220+
# signal and making recency the sole sorting factor. We detect the score range
221+
# and choose the appropriate normalization.
219222
import numpy as np
220223

221-
def sigmoid(x):
224+
def _sigmoid(x: float) -> float:
222225
return 1 / (1 + np.exp(-x))
223226

224-
normalized_scores = [sigmoid(score) for score in scores]
227+
def _rank_normalize_with_ties(score_list: list[float]) -> list[float]:
228+
"""Rank-based normalization that assigns equal ranks to equal scores."""
229+
n = len(score_list)
230+
if n <= 1:
231+
return [1.0] * n
232+
indexed = sorted(enumerate(score_list), key=lambda x: x[1], reverse=True)
233+
result = [0.0] * n
234+
i = 0
235+
while i < n:
236+
j = i
237+
while j < n and indexed[j][1] == indexed[i][1]:
238+
j += 1
239+
# Average rank for tied scores
240+
avg_rank = (i + j - 1) / 2.0
241+
norm = 1.0 - (avg_rank / (n - 1))
242+
for k in range(i, j):
243+
result[indexed[k][0]] = norm
244+
i = j
245+
return result
246+
247+
if scores and min(scores) >= 0.0 and max(scores) <= 1.0:
248+
# Scores are already in [0, 1] (e.g. SiliconFlow, Cohere relevance_score).
249+
# Use rank-based normalization to preserve relative ordering without
250+
# depending on absolute score magnitudes.
251+
normalized_scores = _rank_normalize_with_ties(scores)
252+
else:
253+
# Scores are logits (e.g. local sentence-transformers models).
254+
# Sigmoid maps (-inf, +inf) to (0, 1).
255+
normalized_scores = [_sigmoid(score) for score in scores]
225256

226257
# Create ScoredResult objects with cross-encoder scores
227258
scored_results = []
Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
"""
2+
Unit tests for CrossEncoderReranker.rerank() score normalization logic.
3+
4+
Covers:
5+
1. Rank-based normalization when scores are already in [0, 1].
6+
2. Tied scores receiving identical normalized values.
7+
3. Sigmoid normalization when scores are logits outside [0, 1].
8+
4. Empty candidates returning an empty list without calling predict().
9+
"""
10+
11+
from __future__ import annotations
12+
13+
from unittest.mock import AsyncMock
14+
from uuid import uuid4
15+
16+
import pytest
17+
18+
from hindsight_api.engine.search.reranking import CrossEncoderReranker
19+
from hindsight_api.engine.search.types import MergedCandidate, RetrievalResult
20+
21+
22+
# ---------------------------------------------------------------------------
23+
# Helpers
24+
# ---------------------------------------------------------------------------
25+
26+
def _make_candidates(n: int) -> list[MergedCandidate]:
27+
"""Create *n* minimal MergedCandidate objects."""
28+
candidates = []
29+
for i in range(n):
30+
retrieval = RetrievalResult(
31+
id=str(uuid4()),
32+
text=f"Document {i}",
33+
fact_type="world",
34+
occurred_start=None,
35+
occurred_end=None,
36+
)
37+
candidates.append(
38+
MergedCandidate(retrieval=retrieval, rrf_score=1.0 / (i + 1))
39+
)
40+
return candidates
41+
42+
43+
def _make_cross_encoder(predict_return: list[float]):
44+
"""Return a fake cross-encoder whose `predict` is an AsyncMock."""
45+
ce = AsyncMock()
46+
ce.predict = AsyncMock(return_value=predict_return)
47+
ce.provider_name = "local"
48+
ce.initialize = AsyncMock()
49+
return ce
50+
51+
52+
# ---------------------------------------------------------------------------
53+
# Tests
54+
# ---------------------------------------------------------------------------
55+
56+
@pytest.mark.asyncio
57+
async def test_rank_normalization_for_0_1_scores():
58+
"""Scores already in [0, 1] should use rank-based normalization."""
59+
# Scores in [0, 1] but NOT logit-shaped — rank normalization preserves order
60+
raw_scores = [0.1, 0.5, 0.9]
61+
ce = _make_cross_encoder(raw_scores)
62+
reranker = CrossEncoderReranker(cross_encoder=ce)
63+
reranker._initialized = True
64+
65+
candidates = _make_candidates(3)
66+
results = await reranker.rerank("test query", candidates)
67+
68+
assert len(results) == 3
69+
# Top-ranked (0.9) gets normalized 1.0
70+
assert results[0].cross_encoder_score == pytest.approx(0.9)
71+
assert results[0].cross_encoder_score_normalized == pytest.approx(1.0)
72+
# Bottom-ranked (0.1) gets normalized 0.0
73+
assert results[-1].cross_encoder_score == pytest.approx(0.1)
74+
assert results[-1].cross_encoder_score_normalized == pytest.approx(0.0)
75+
# Middle gets 0.5
76+
assert results[1].cross_encoder_score_normalized == pytest.approx(0.5)
77+
# predict was called exactly once
78+
ce.predict.assert_awaited_once()
79+
80+
81+
@pytest.mark.asyncio
82+
async def test_tied_scores_get_same_normalized_value():
83+
"""When raw scores contain ties, tied entries must receive the same normalized score."""
84+
# Two scores tied at 0.7, one distinct at 0.3
85+
raw_scores = [0.7, 0.3, 0.7]
86+
ce = _make_cross_encoder(raw_scores)
87+
reranker = CrossEncoderReranker(cross_encoder=ce)
88+
reranker._initialized = True
89+
90+
candidates = _make_candidates(3)
91+
results = await reranker.rerank("test query", candidates)
92+
93+
# The two 0.7 entries should share the same normalized value.
94+
# With rank-based normalization over 3 items (indices 0,1,2 sorted desc):
95+
# rank 0 & 1 tied → avg_rank = 0.5 → norm = 1 - 0.5/2 = 0.75
96+
# rank 2 → norm = 1 - 2/2 = 0.0
97+
# After sorting by normalized score descending, the two 0.7 entries are
98+
# at indices 0 and 1; the 0.3 entry is at index 2.
99+
scores_by_raw = {r.cross_encoder_score: r.cross_encoder_score_normalized for r in results}
100+
tied_norm = scores_by_raw[0.7]
101+
assert tied_norm == pytest.approx(0.75)
102+
assert scores_by_raw[0.3] == pytest.approx(0.0)
103+
# Both 0.7 results have identical normalized scores
104+
assert results[0].cross_encoder_score_normalized == pytest.approx(
105+
results[1].cross_encoder_score_normalized
106+
)
107+
108+
109+
@pytest.mark.asyncio
110+
async def test_sigmoid_normalization_for_logits():
111+
"""When scores are outside [0, 1] (logits), sigmoid normalization is used."""
112+
raw_scores = [2.0, -1.0, 0.0]
113+
ce = _make_cross_encoder(raw_scores)
114+
reranker = CrossEncoderReranker(cross_encoder=ce)
115+
reranker._initialized = True
116+
117+
candidates = _make_candidates(3)
118+
results = await reranker.rerank("test query", candidates)
119+
120+
assert len(results) == 3
121+
import math
122+
123+
# Results are sorted by weight descending
124+
expected_sigmoid = [1 / (1 + math.exp(-s)) for s in raw_scores]
125+
expected_sorted = sorted(expected_sigmoid, reverse=True)
126+
127+
for result, expected in zip(results, expected_sorted):
128+
assert result.cross_encoder_score_normalized == pytest.approx(expected, rel=1e-6)
129+
130+
# Verify the highest logit (2.0) maps to the highest normalized score
131+
assert results[0].cross_encoder_score == pytest.approx(2.0)
132+
assert results[0].cross_encoder_score_normalized > 0.5
133+
134+
135+
@pytest.mark.asyncio
136+
async def test_empty_candidates_returns_empty_without_predict():
137+
"""When candidates are empty, rerank must return [] without calling predict()."""
138+
ce = _make_cross_encoder([])
139+
reranker = CrossEncoderReranker(cross_encoder=ce)
140+
reranker._initialized = True
141+
142+
results = await reranker.rerank("test query", [])
143+
144+
assert results == []
145+
ce.predict.assert_not_awaited()
146+
147+
148+
@pytest.mark.asyncio
149+
async def test_mixed_boundary_scores_use_rank():
150+
"""Boundary values 0.0 and 1.0 (still in [0,1]) should trigger rank normalization."""
151+
raw_scores = [0.0, 1.0, 0.5]
152+
ce = _make_cross_encoder(raw_scores)
153+
reranker = CrossEncoderReranker(cross_encoder=ce)
154+
reranker._initialized = True
155+
156+
candidates = _make_candidates(3)
157+
results = await reranker.rerank("test query", candidates)
158+
159+
# Rank-based: 1.0 → 1.0, 0.5 → 0.5, 0.0 → 0.0
160+
by_score = {r.cross_encoder_score: r.cross_encoder_score_normalized for r in results}
161+
assert by_score[1.0] == pytest.approx(1.0)
162+
assert by_score[0.5] == pytest.approx(0.5)
163+
assert by_score[0.0] == pytest.approx(0.0)

0 commit comments

Comments
 (0)