diff --git a/axon/services/reranker.py b/axon/services/reranker.py new file mode 100644 index 0000000..dcdc1c2 --- /dev/null +++ b/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)