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

@@ -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)