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)