Merge pull request #529 from ggozad/fix/cross-encoder-activation

Keep cross-encoder rerank scores apart when they saturate

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

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

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
+        ]

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` "

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

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)