fix(rerank): normalize reranker scores onto a single scale before hybrid blend (#15429)

### What problem does this PR solve?

Closes #15428

The hybrid score in `rag/nlp/search.py` (`rerank_by_model`) blends
reranker similarity with token similarity on a fixed `[0, 1]` scale:

```python
return tkweight * np.array(tksim) + vtweight * vtsim + rank_fea  # tkweight=0.3, vtweight=0.7
```

The reranker implementations did not agree on that scale. Only three of
roughly 17 providers normalized their output, and `NvidiaRerank`
returned raw, unbounded logits. Weighted at `0.7`, a negative logit
could push a genuinely relevant chunk below pure keyword matches, and
its magnitude swamped `tksim`, which lives in `[0, 1]`. The practical
effect was that the same query produced differently scaled scores
depending on the configured reranker, and logit based providers degraded
retrieval quality instead of improving it.

This PR enforces a single scoring contract in one place:

- `Base.similarity` is now the only public entry point. It
short-circuits empty input and guarantees a normalized result. Each
provider implements its raw scoring in `_compute_rank`, which removes
sixteen duplicated empty input guards and the three scattered
normalization calls.
- Normalization is range aware. Providers that already return calibrated
`[0, 1]` relevance scores (Cohere, Jina, Voyage, and others) keep their
absolute magnitudes, so `similarity_threshold` filtering and the
reported `vector_similarity` stay meaningful. Only out-of-range output
such as NVIDIA logits is min-max rescaled into `[0, 1]`.
- The twelve leftover `[DEBUG ...]` prints in `rerank_by_model`,
introduced in #14231, are removed. They ran on every retrieval, added
per chunk overhead, and leaked queries, keywords, and document content
to stdout and logs.

A new regression suite in
`test/unit_test/rag/llm/test_rerank_normalization.py` covers logit
rescaling (positive, negative, and flat batches), preservation of
already calibrated scores, ordering, empty input handling, and the per
provider HTTP path. It also asserts that no provider overrides
`similarity()`, so the contract cannot silently drift.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
cleanjunc
2026-06-08 06:53:22 +03:00
committed by GitHub
parent 3d7adf2193
commit 38f9ea5fec
3 changed files with 242 additions and 80 deletions

View File

@@ -32,21 +32,63 @@ class Base(ABC):
pass
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
raise NotImplementedError("Please implement encode method!")
"""Score ``texts`` against ``query`` and return ``(rank, token_count)``.
This is the single public entry point shared by every reranker. It
short-circuits empty input and guarantees the returned scores are
min-max normalized to ``[0, 1]`` regardless of what the backend emits
(relevance scores, cosine similarities or raw logits). Downstream
hybrid scoring blends the reranker output with token similarity on a
fixed ``[0, 1]`` scale, so an un-normalized provider (e.g. NVIDIA's
unbounded logits) would otherwise corrupt the final ordering.
Subclasses implement provider-specific scoring in :meth:`_compute_rank`
and must not normalize themselves.
"""
if not query or not texts:
return np.zeros(len(texts) if texts else 0, dtype=float), 0
rank, token_count = self._compute_rank(query, texts)
rank = np.asarray(rank, dtype=float)
if rank.size:
logging.debug(
"Rerank %s scores before normalization: count=%d min=%.4f max=%.4f",
self.__class__.__name__,
rank.size,
float(np.min(rank)),
float(np.max(rank)),
)
return self._normalize_rank(rank), token_count
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
"""Provider-specific scoring. ``query`` and ``texts`` are non-empty."""
raise NotImplementedError("Please implement _compute_rank method!")
@staticmethod
def _normalize_rank(rank: np.ndarray) -> np.ndarray:
"""Guarantee scores land in ``[0, 1]`` for the hybrid blend.
Providers that already emit calibrated relevance scores in ``[0, 1]``
(Cohere, Jina, Voyage, ...) are returned unchanged, so their absolute
magnitudes, ``similarity_threshold`` semantics and reported
``vector_similarity`` are preserved. Only out-of-range output (e.g.
NVIDIA's unbounded, often negative logits) is rescaled: a batch with a
usable spread is min-max mapped onto ``[0, 1]`` (which stops a negative
logit from dragging a relevant chunk below pure keyword matches once
weighted by ``vtweight``), while a spreadless batch (including a single
candidate) has no relative signal and is clamped instead, so a lone
high score is not silently zeroed.
"""
if rank.size == 0:
return rank
min_rank = np.min(rank)
max_rank = np.max(rank)
min_rank = float(np.min(rank))
max_rank = float(np.max(rank))
if not np.isclose(min_rank, max_rank, atol=1e-3):
rank = (rank - min_rank) / (max_rank - min_rank)
else:
rank = np.zeros_like(rank)
return rank
if min_rank >= 0.0 and max_rank <= 1.0:
return rank
span = max_rank - min_rank
if span < 1e-3:
return np.clip(rank, 0.0, 1.0)
return (rank - min_rank) / span
class JinaRerank(Base):
@@ -57,9 +99,7 @@ class JinaRerank(Base):
self.headers = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}
self.model_name = model_name
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts) if texts else 0, dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
texts = [truncate(t, 8196) for t in texts]
data = {"model": self.model_name, "query": query, "documents": texts, "top_n": len(texts)}
response = requests.post(self.base_url, headers=self.headers, json=data, timeout=30)
@@ -88,9 +128,7 @@ class XInferenceRerank(Base):
if key and key != "x":
self.headers["Authorization"] = f"Bearer {key}"
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts) if texts else 0, dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
pairs = [(query, truncate(t, 4096)) for t in texts]
token_count = 0
for _, t in pairs:
@@ -119,9 +157,7 @@ class LocalAIRerank(Base):
self.headers = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}
self.model_name = model_name.split("___")[0]
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
texts = [truncate(t, 500) for t in texts]
data = {
"model": self.model_name,
@@ -141,8 +177,6 @@ class LocalAIRerank(Base):
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)
rank = Base._normalize_rank(rank)
return rank, token_count
@@ -167,9 +201,7 @@ class NvidiaRerank(Base):
"Authorization": f"Bearer {key}",
}
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
token_count = num_tokens_from_string(query) + sum([num_tokens_from_string(t) for t in texts])
data = {
"model": self.model_name,
@@ -196,7 +228,7 @@ class LmStudioRerank(Base):
def __init__(self, key, model_name, base_url, **kwargs):
pass
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
raise NotImplementedError("The LmStudioRerank has not been implemented")
@@ -212,9 +244,7 @@ class OpenAI_APIRerank(Base):
self.headers = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}
self.model_name = model_name.split("___")[0]
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
texts = [truncate(t, 500) for t in texts]
data = {
"model": self.model_name,
@@ -234,8 +264,6 @@ class OpenAI_APIRerank(Base):
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)
rank = Base._normalize_rank(rank)
return rank, token_count
@@ -251,9 +279,7 @@ class CoHereRerank(Base):
self.client = Client(**client_kwargs)
self.model_name = model_name.split("___")[0]
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
token_count = num_tokens_from_string(query) + sum([num_tokens_from_string(t) for t in texts])
res = self.client.rerank(
model=self.model_name,
@@ -277,7 +303,7 @@ class TogetherAIRerank(Base):
def __init__(self, key, model_name, base_url, **kwargs):
pass
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
raise NotImplementedError("The api has not been implemented")
@@ -298,9 +324,7 @@ class SILICONFLOWRerank(Base):
"authorization": f"Bearer {key}",
}
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
payload = {
"model": self.model_name,
"query": query,
@@ -334,9 +358,7 @@ class BaiduYiyanRerank(Base):
self.client = Reranker(ak=ak, sk=sk, request_timeout=30)
self.model_name = model_name
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
res = self.client.do(
model=self.model_name,
query=query,
@@ -361,9 +383,7 @@ class VoyageRerank(Base):
self.client = voyageai.Client(api_key=key, timeout=30.0)
self.model_name = model_name
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts) if texts else 0, dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
rank = np.zeros(len(texts), dtype=float)
res = self.client.rerank(query=query, documents=texts, model=self.model_name, top_k=len(texts))
@@ -385,10 +405,7 @@ class QWenRerank(Base):
# Remove invalid global timeout, use official SDK per-request timeout parameter
self.request_timeout = 30.0
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
import dashscope
# Pass official request_timeout parameter to both API call branches
@@ -455,9 +472,7 @@ class HuggingfaceRerank(Base):
self.model_name = model_name.split("___")[0]
self.base_url = base_url
def similarity(self, query: str, texts: List) -> tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> tuple[np.ndarray, int]:
token_count = 0
for t in texts:
token_count += num_tokens_from_string(t)
@@ -479,10 +494,7 @@ class GPUStackRerank(Base):
"authorization": f"Bearer {key}",
}
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
payload = {
"model": self.model_name,
"query": query,
@@ -535,9 +547,7 @@ class Ai302Rerank(Base):
self.headers = {"Content-Type": "application/json", "Authorization": f"Bearer {key}"}
self.model_name = model_name
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
texts = [truncate(t, 500) for t in texts]
data = {"model": self.model_name, "query": query, "documents": texts, "top_n": len(texts)}
response = requests.post(self.base_url, headers=self.headers, json=data, timeout=30)
@@ -585,10 +595,7 @@ class RAGconRerank(Base):
self.model_name = model_name
def similarity(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
if not query or not texts:
return np.zeros(len(texts), dtype=float), 0
def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]:
texts = [truncate(t, 500) for t in texts]
data = {
"model": self.model_name,
@@ -606,6 +613,4 @@ class RAGconRerank(Base):
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)
rank = Base._normalize_rank(rank)
return rank, token_count

View File

@@ -513,9 +513,7 @@ class Dealer:
def rerank_by_model(self, rerank_mdl, sres, query, tkweight=0.3,
vtweight=0.7, cfield="content_ltks",
rank_feature: dict | None = None):
print(f"[DEBUG rerank_by_model] query={query}, tkweight={tkweight}, vtweight={vtweight}")
_, keywords = self.qryr.question(query)
print(f"[DEBUG rerank_by_model] keywords={keywords}")
for i in sres.ids:
if isinstance(sres.field[i].get("important_kwd", []), str):
@@ -528,29 +526,16 @@ class Dealer:
important_kwd = sres.field[i].get("important_kwd", [])
tks = content_ltks + title_tks + important_kwd
ins_tw.append(tks)
print(f"[DEBUG rerank_by_model] chunk id={i}, content_ltks={len(content_ltks)}, title_tks={len(title_tks)}, important_kwd={len(important_kwd)}")
doc_text = remove_redundant_spaces(" ".join(tks))
if len(doc_text) > 100:
print(f"[DEBUG rerank_by_model] chunk id={i}, doc_text (first 100)={doc_text[:100]}...")
else:
print(f"[DEBUG rerank_by_model] chunk id={i}, doc_text={doc_text}")
docs = [remove_redundant_spaces(" ".join(tks)) for tks in ins_tw]
print(f"[DEBUG rerank_by_model] docs sent to reranker: {len(docs)} docs")
for idx, doc in enumerate(docs[:2]): # Print first 2
print(f"[DEBUG rerank_by_model] doc[{idx}] len={len(doc)}, full={doc}")
if len(doc) > 100:
print(f"[DEBUG rerank_by_model] doc[{idx}] (first 100)={doc[:100]}...")
else:
print(f"[DEBUG rerank_by_model] doc[{idx}]={doc}")
tksim = self.qryr.token_similarity(keywords, ins_tw)
print(f"[DEBUG rerank_by_model] tksim={tksim}")
# rerank_mdl.similarity() returns scores normalized to [0, 1] for every
# provider (see RerankModel.Base.similarity), so the blend below stays
# on a single scale regardless of the configured reranker.
vtsim, _ = rerank_mdl.similarity(query, docs)
print(f"[DEBUG rerank_by_model] vtsim from reranker={vtsim}")
## For rank feature(tag_fea) scores.
rank_fea = self._rank_feature_scores(rank_feature, sres)
print(f"[DEBUG rerank_by_model] rank_fea={rank_fea}")
return tkweight * np.array(tksim) + vtweight * vtsim + rank_fea, tksim, vtsim

View File

@@ -0,0 +1,172 @@
#
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Regression tests for the shared reranker score-normalization contract.
Every reranker must return scores on a single ``[0, 1]`` scale so that the
hybrid blend in ``rag/nlp/search.py`` (``tkweight * tksim + vtweight * vtsim``)
stays comparable across providers. Historically only 3 of ~17 providers
normalized, and NVIDIA returned raw, unbounded logits — which corrupted
retrieval ordering. The contract is now enforced once in ``Base.similarity``.
"""
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
from rag.llm.rerank_model import (
Base,
JinaRerank,
NvidiaRerank,
)
pytestmark = pytest.mark.p1
def _mock_post(payload):
"""Patch ``requests.post`` so ``response.json()`` returns ``payload``."""
response = MagicMock()
response.raise_for_status.return_value = None
response.json.return_value = payload
return patch("rag.llm.rerank_model.requests.post", return_value=response)
class _RawRerank(Base):
"""Minimal provider that emits arbitrary raw scores via ``_compute_rank``."""
def __init__(self, raw):
self._raw = np.asarray(raw, dtype=float)
def _compute_rank(self, query, texts):
return self._raw, 0
# --- The central guarantee: every provider's output lands in [0, 1] ----------
@pytest.mark.parametrize(
"raw, expected",
[
# Unbounded NVIDIA-style logits, including negatives -> rescaled.
([10.0, -3.0, 0.0], [1.0, 0.0, 3.0 / 13.0]),
# Large positive logits -> rescaled.
([100.0, 50.0, 75.0], [1.0, 0.0, 0.5]),
# Negative-only logits -> rescaled.
([-1.0, -5.0, -3.0], [1.0, 0.0, 0.5]),
],
)
def test_out_of_range_scores_are_rescaled(raw, expected):
rank, _ = _RawRerank(raw).similarity("q", ["a", "b", "c"])
assert np.allclose(rank, expected)
assert rank.min() >= 0.0 and rank.max() <= 1.0
@pytest.mark.parametrize(
"raw",
[
[0.9, 0.1, 0.5], # spread relevance scores
[0.8, 0.8, 0.8], # all-equal but valid -> not zeroed
[1.0], # single calibrated candidate -> not zeroed
[0.0, 1.0, 0.42], # already spanning the full range
],
)
def test_in_range_scores_are_preserved(raw):
# Calibrated [0,1] providers (Cohere/Jina/Voyage/...) keep their absolute
# magnitudes, so similarity_threshold and reported vector_similarity stay
# meaningful and degenerate batches are NOT collapsed to zero.
rank, _ = _RawRerank(raw).similarity("q", ["x"] * len(raw))
assert np.allclose(rank, raw)
def test_normalization_preserves_ordering():
raw = [-5.0, 12.0, 3.0, -1.0]
rank, _ = _RawRerank(raw).similarity("q", ["a", "b", "c", "d"])
assert list(np.argsort(rank)) == list(np.argsort(raw))
@pytest.mark.parametrize(
"raw, expected",
[
# Single out-of-range candidate: clamped, never zeroed and never NaN.
([5.0], [1.0]),
([-3.0], [0.0]),
# Spreadless out-of-range batch: clamped per element, not collapsed.
([5.0, 5.0, 5.0], [1.0, 1.0, 1.0]),
([-2.0, -2.0, -2.0], [0.0, 0.0, 0.0]),
],
)
def test_spreadless_out_of_range_batch_is_clamped(raw, expected):
rank, _ = _RawRerank(raw).similarity("q", ["x"] * len(raw))
assert np.allclose(rank, expected)
assert not np.isnan(rank).any()
# --- Empty input short-circuits before any backend call ----------------------
@pytest.mark.parametrize("query, texts", [("", ["a"]), ("q", []), ("", [])])
def test_empty_input_returns_zeros_without_backend(query, texts):
provider = _RawRerank([1.0])
provider._compute_rank = MagicMock(side_effect=AssertionError("backend called"))
rank, tokens = provider.similarity(query, texts)
assert tokens == 0
assert rank.size == len(texts)
assert rank.dtype == float
# --- Per-provider: raw backend payloads come out normalized ------------------
def test_nvidia_logits_are_normalized():
"""NVIDIA emits raw logits; without central normalization a negative logit
with vtweight=0.7 would sink a relevant chunk below keyword matches."""
nv = NvidiaRerank("key", "nvidia/rerank-qa-mistral-4b")
payload = {"rankings": [{"index": 0, "logit": 8.0}, {"index": 1, "logit": -4.0}, {"index": 2, "logit": 1.0}]}
with _mock_post(payload):
rank, _ = nv.similarity("q", ["a", "b", "c"])
# _compute_rank still returns the raw logits (no per-provider normalization)...
with _mock_post(payload):
raw, _ = nv._compute_rank("q", ["a", "b", "c"])
assert raw.min() < 0 # genuinely unbounded/negative
# ...but the public contract normalizes them.
assert np.allclose(rank, [1.0, 0.0, 5.0 / 12.0])
assert rank.min() >= 0.0 and rank.max() <= 1.0
def test_calibrated_relevance_scores_are_preserved():
# A provider already returning [0,1] relevance scores keeps them verbatim;
# min-max would have stretched these to [1.0, 0.0, 0.5].
jina = JinaRerank("key", base_url="http://x/rerank")
payload = {"results": [{"index": 0, "relevance_score": 0.8}, {"index": 1, "relevance_score": 0.2}, {"index": 2, "relevance_score": 0.5}]}
with _mock_post(payload):
rank, _ = jina.similarity("q", ["a", "b", "c"])
assert np.allclose(rank, [0.8, 0.2, 0.5])
# --- Structural guarantee: providers override _compute_rank, not similarity --
def test_providers_share_single_similarity_entrypoint():
import inspect
import rag.llm.rerank_model as rm
overrides = []
for _, cls in inspect.getmembers(rm, inspect.isclass):
if issubclass(cls, Base) and cls is not Base and "similarity" in cls.__dict__:
overrides.append(cls.__name__)
assert overrides == [], f"providers must not override similarity(): {overrides}"