Import source: axon/services/reranker

Import the existing local project snapshot. This commit organizes existing files; it does not represent newly developed functionality.

axon/services/reranker.py

@@ -0,0 +1,75 @@
+import asyncio
+import logging
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+_model = None
+_predict_lock = asyncio.Lock()
+
+
+def load_reranker() -> None:
+    global _model
+    from axon.config import settings
+    from sentence_transformers import CrossEncoder
+
+    logger.info("Loading reranker model: %s", settings.reranker_model)
+    _model = CrossEncoder(settings.reranker_model)
+    logger.info("Reranker model loaded")
+
+
+def get_reranker() -> "CrossEncoder | None":
+    return _model
+
+
+def chunk_text(text: str, max_tokens: int = 512, max_chunks: int = 4) -> list[str]:
+    if not text:
+        return [""]
+    words = text.split()
+    words_per_chunk = max(1, int(max_tokens * 0.75))
+    chunks = [
+        " ".join(words[i : i + words_per_chunk])
+        for i in range(0, len(words), words_per_chunk)
+    ]
+    return (chunks or [""])[:max_chunks]
+
+
+def _sigmoid(x: np.ndarray) -> np.ndarray:
+    return 1.0 / (1.0 + np.exp(-x))
+
+
+async def rerank_results(
+    query: str,
+    contents: list[tuple[str, str]],
+) -> list[tuple[str, str, float]]:
+    if not contents:
+        return []
+
+    # Build (url, chunks) pairs and flat passage list
+    url_chunks: list[tuple[str, list[str]]] = []
+    flat_passages: list[str] = []
+    for url, text in contents:
+        chunks = chunk_text(text)
+        url_chunks.append((url, chunks))
+        flat_passages.extend(chunks)
+
+    pairs = [[query, p] for p in flat_passages]
+
+    def _predict():
+        return _model.predict(pairs)
+
+    async with _predict_lock:
+        raw_scores = await asyncio.to_thread(_predict)
+    scores: list[float] = _sigmoid(np.array(raw_scores, dtype=float)).tolist()
+
+    # Pick best chunk per URL
+    results: list[tuple[str, str, float]] = []
+    idx = 0
+    for url, chunks in url_chunks:
+        chunk_scores = scores[idx : idx + len(chunks)]
+        best_i = int(np.argmax(chunk_scores))
+        results.append((url, chunks[best_i], round(chunk_scores[best_i], 6)))
+        idx += len(chunks)
+
+    return sorted(results, key=lambda t: t[2], reverse=True)