feat: cross-encoder reranker with ONNX inference

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

src/mediastorm/rag/reranker.py

@@ -0,0 +1,47 @@
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+from tokenizers import Tokenizer
+
+from mediastorm.config import RERANKER_MODEL_DIR
+
+
+class Reranker:
+    def __init__(self, model_dir: Path | str = RERANKER_MODEL_DIR):
+        model_dir = Path(model_dir)
+        self._tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
+        self._tokenizer.enable_padding()
+        self._tokenizer.enable_truncation(max_length=512)
+        self._session = ort.InferenceSession(
+            str(model_dir / "model.onnx"),
+            providers=["CPUExecutionProvider"],
+        )
+
+    def score_pairs(self, query: str, documents: list[str]) -> list[float]:
+        """Score (query, document) pairs. Returns one relevance score per document."""
+        if not documents:
+            return []
+
+        encoded = self._tokenizer.encode_batch(
+            [(query, doc) for doc in documents],
+        )
+
+        input_ids = np.array([e.ids for e in encoded], dtype=np.int64)
+        attention_mask = np.array([e.attention_mask for e in encoded], dtype=np.int64)
+
+        feeds = {"input_ids": input_ids, "attention_mask": attention_mask}
+
+        # Add token_type_ids if the model expects them
+        input_names = [inp.name for inp in self._session.get_inputs()]
+        if "token_type_ids" in input_names:
+            token_type_ids = np.array([e.type_ids for e in encoded], dtype=np.int64)
+            feeds["token_type_ids"] = token_type_ids
+
+        outputs = self._session.run(None, feeds)
+        logits = outputs[0]
+
+        # Cross-encoder output: (batch_size, 1) logit — squeeze to flat list
+        if logits.ndim == 2 and logits.shape[1] == 1:
+            return logits[:, 0].tolist()
+        return logits.flatten().tolist()

tests/test_reranker.py

@@ -0,0 +1,39 @@
+import pytest
+from mediastorm.rag.reranker import Reranker
+
+
+@pytest.fixture(scope="module")
+def reranker():
+    return Reranker()
+
+
+def test_score_pairs_returns_one_score_per_document(reranker):
+    scores = reranker.score_pairs("climate change", ["global warming effects", "jazz music history"])
+    assert len(scores) == 2
+    assert all(isinstance(s, float) for s in scores)
+
+
+def test_relevant_document_scores_higher(reranker):
+    scores = reranker.score_pairs(
+        "climate change",
+        ["global warming and rising sea levels", "jazz musicians in New York"],
+    )
+    assert scores[0] > scores[1], f"Relevant doc should score higher: {scores}"
+
+
+def test_score_pairs_deterministic(reranker):
+    s1 = reranker.score_pairs("test query", ["test document"])
+    s2 = reranker.score_pairs("test query", ["test document"])
+    assert s1[0] == pytest.approx(s2[0], abs=1e-5)
+
+
+def test_score_pairs_handles_long_text(reranker):
+    long_text = "word " * 1000
+    scores = reranker.score_pairs("query", [long_text])
+    assert len(scores) == 1
+    assert isinstance(scores[0], float)
+
+
+def test_score_pairs_empty_list(reranker):
+    scores = reranker.score_pairs("query", [])
+    assert scores == []