mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-08 12:24:48 +08:00
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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
172
test/unit_test/rag/llm/test_rerank_normalization.py
Normal file
172
test/unit_test/rag/llm/test_rerank_normalization.py
Normal 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}"
|
||||
Reference in New Issue
Block a user