From d63199d96d7a02914ec336e381eef538b7452ead Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 7 Aug 2026 13:35:38 +0300 Subject: [PATCH] Keep cross-encoder rerank scores apart when they saturate mxbai-rerank-base-v2 ships a Sigmoid activation and evaluates it in bf16, so every strongly-relevant candidate rounds to exactly 1.0. Ties then leave the order to the stable sort, which preserves the incoming hybrid ranking: on 100 t2_finqa retrieval cases the reranker scored MAP 0.661 against 0.659 with no reranker at all, and 0.742 once the scores separate. Ask the model for logits and apply the sigmoid here, where it runs in float64. Scores stay 0-1, matching the cohere, vllm and zeroentropy rerankers. Also drop the remaining pyright references; the project type-checks with ty. --- CHANGELOG.md | 4 ++ app/backend/pyproject.toml | 2 +- .../haiku/rag/reranking/cross_encoder.py | 18 +++++-- .../haiku/rag/reranking/jina_local.py | 4 +- tests/json_body_serializer.py | 1 - tests/test_reranker.py | 47 ++++++++++++++++++- 6 files changed, 64 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 81e6756d..fa48ed23 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Fixed + +- `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1. + ## [0.73.0] - 2026-08-06 ### Added diff --git a/app/backend/pyproject.toml b/app/backend/pyproject.toml index 7164c9fe..dcb00578 100644 --- a/app/backend/pyproject.toml +++ b/app/backend/pyproject.toml @@ -13,7 +13,7 @@ dependencies = [ ] [dependency-groups] -dev = ["pyright>=1.1.407", "ruff>=0.14.10"] +dev = ["ty>=0.0.28", "ruff>=0.14.10"] [tool.hatch.metadata] allow-direct-references = true diff --git a/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py b/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py index a12d45ee..a27f3954 100644 --- a/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py +++ b/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py @@ -1,9 +1,9 @@ import asyncio +import math try: - from sentence_transformers import ( - CrossEncoder, # pyright: ignore[reportMissingImports] - ) + import torch + from sentence_transformers import CrossEncoder except ImportError as e: # pragma: no cover raise ImportError( "sentence-transformers is not installed. Install it with " @@ -30,7 +30,15 @@ class CrossEncoderReranker(RerankerBase): self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: documents = [chunk.content for chunk in chunks] + # Ask for logits and squash them here: the model's own sigmoid runs in + # bf16, where saturated scores round onto identical values and leave the + # order of the top candidates to the sort. rankings = await asyncio.to_thread( - lambda: self._reranker.rank(query, documents, top_k=top_n) + lambda: self._reranker.rank( + query, documents, top_k=top_n, activation_fn=torch.nn.Identity() + ) ) - return [(chunks[r["corpus_id"]], float(r["score"])) for r in rankings] + return [ + (chunks[r["corpus_id"]], 1.0 / (1.0 + math.exp(-r["score"]))) + for r in rankings + ] diff --git a/haiku_rag_slim/haiku/rag/reranking/jina_local.py b/haiku_rag_slim/haiku/rag/reranking/jina_local.py index 16bc463c..39a7ccee 100644 --- a/haiku_rag_slim/haiku/rag/reranking/jina_local.py +++ b/haiku_rag_slim/haiku/rag/reranking/jina_local.py @@ -1,9 +1,7 @@ import asyncio try: - from transformers import ( - AutoModel, # pyright: ignore[reportMissingImports] - ) + from transformers import AutoModel except ImportError as e: # pragma: no cover raise ImportError( "transformers is not installed. Please install it with `pip install transformers torch` " diff --git a/tests/json_body_serializer.py b/tests/json_body_serializer.py index 82ef00d3..d460cd6b 100644 --- a/tests/json_body_serializer.py +++ b/tests/json_body_serializer.py @@ -1,5 +1,4 @@ # Adapted from pydantic-ai: https://github.com/pydantic/pydantic-ai/blob/main/tests/json_body_serializer.py -# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false import json import urllib.parse from typing import TYPE_CHECKING, Any diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 635fe5fc..1dea945f 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -481,16 +481,56 @@ async def test_cross_encoder_reranker(): pytest.skip("sentence-transformers not installed") +@pytest.mark.asyncio +@pytest.mark.integration +async def test_cross_encoder_separates_saturated_scores(): + """`mxbai-rerank-base-v2` ships a Sigmoid it evaluates in bf16, where every + strongly-relevant candidate rounds to exactly 1.0. Squashing the logits + ourselves keeps them apart.""" + try: + from haiku.rag.reranking.cross_encoder import CrossEncoderReranker + + saturating = [ + Chunk(content=content, document_id=str(i)) + for i, content in enumerate( + [ + "To Kill a Mockingbird is a novel by Harper Lee published in 1960.", + "Harper Lee wrote To Kill a Mockingbird, published in 1960.", + "The author of To Kill a Mockingbird is Harper Lee.", + "To Kill a Mockingbird, written by Harper Lee, appeared in 1960.", + "Harper Lee, author of To Kill a Mockingbird, won a Pulitzer Prize.", + "Harper Lee is the novelist who wrote To Kill a Mockingbird.", + ] + ) + ] + + reranker = CrossEncoderReranker("mixedbread-ai/mxbai-rerank-base-v2") + reranked = await reranker.rerank( + "Who wrote 'To Kill a Mockingbird'?", saturating, top_n=len(saturating) + ) + + scores = [score for chunk, score in reranked] + assert len(set(scores)) == len(scores) + assert all(0.0 < score < 1.0 for score in scores) + except ImportError: + pytest.skip("sentence-transformers not installed") + + @pytest.mark.asyncio async def test_cross_encoder_reranks_via_model_ranking(monkeypatch): """The rank() results map back onto the input chunks by corpus_id.""" + import math + + import torch + from haiku.rag.reranking import cross_encoder as ce_module class _StubCrossEncoder: def __init__(self, model): self.model = model - def rank(self, query, documents, top_k=10): + def rank(self, query, documents, top_k=10, activation_fn=None): + self.activation_fn = activation_fn # Reverse order so the mapping back to chunks is observable. return [ {"corpus_id": i, "score": 1.0 - (i / 10)} @@ -505,4 +545,7 @@ async def test_cross_encoder_reranks_via_model_ranking(monkeypatch): assert len(reranked) == 2 last_index = len(chunks) - 1 assert reranked[0][0] is chunks[last_index] - assert reranked[0][1] == pytest.approx(1.0 - last_index / 10) + # The stub returns logits; the reranker squashes them itself. + stub_logit = 1.0 - last_index / 10 + assert reranked[0][1] == pytest.approx(1.0 / (1.0 + math.exp(-stub_logit))) + assert isinstance(reranker._reranker.activation_fn, torch.nn.Identity)