diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index 99801e00a..c149b2307 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -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 diff --git a/rag/nlp/search.py b/rag/nlp/search.py index 8201bb370..3a702f3bc 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -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 diff --git a/test/unit_test/rag/llm/test_rerank_normalization.py b/test/unit_test/rag/llm/test_rerank_normalization.py new file mode 100644 index 000000000..5af9abc46 --- /dev/null +++ b/test/unit_test/rag/llm/test_rerank_normalization.py @@ -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}"