Spaces:
Sleeping
Sleeping
| """ | |
| Cross-Encoder Reranking Service | |
| After initial vector/BM25 retrieval returns candidate chunks, | |
| a cross-encoder model reads the query + each chunk together | |
| and produces a true semantic relevance score. | |
| This catches cases where cosine similarity ranks a less-relevant | |
| chunk higher than a more-relevant one. | |
| Model: cross-encoder/ms-marco-MiniLM-L-6-v2 | |
| - Trained on MS MARCO passage ranking | |
| - ~80M params, runs on CPU in ~50ms per batch of 20 | |
| - Free, no API key needed | |
| """ | |
| import logging | |
| from typing import List | |
| from app.config import settings | |
| logger = logging.getLogger(__name__) | |
| _model = None | |
| def _get_model(): | |
| """Lazy-load the cross-encoder model (loaded once, cached globally).""" | |
| global _model | |
| if _model is None: | |
| try: | |
| from sentence_transformers import CrossEncoder | |
| logger.info(f"Loading reranker model: {settings.RERANK_MODEL}") | |
| _model = CrossEncoder(settings.RERANK_MODEL) | |
| logger.info("Reranker model loaded successfully") | |
| except ImportError: | |
| logger.error( | |
| "sentence-transformers not installed. " | |
| "Install with: pip install sentence-transformers" | |
| ) | |
| raise | |
| return _model | |
| def rerank( | |
| query: str, | |
| chunks: List[dict], | |
| top_k: int = None, | |
| text_key: str = "text", | |
| ) -> List[dict]: | |
| """ | |
| Rerank chunks by cross-encoder relevance score. | |
| Args: | |
| query: The user's search query | |
| chunks: List of chunk dicts, each must have `text_key` field | |
| top_k: Number of top chunks to return (default: settings.RERANK_TOP_K) | |
| text_key: Key in chunk dict containing the text to score | |
| Returns: | |
| Top-k chunks sorted by cross-encoder score (descending). | |
| Each chunk gets an added `rerank_score` field. | |
| """ | |
| if not chunks: | |
| return [] | |
| if top_k is None: | |
| top_k = settings.RERANK_TOP_K | |
| # If only 1 chunk, no reranking needed | |
| if len(chunks) <= top_k: | |
| for c in chunks: | |
| c["rerank_score"] = c.get("score", 0.0) | |
| return chunks | |
| model = _get_model() | |
| # Build query-document pairs for cross-encoder | |
| pairs = [(query, chunk[text_key]) for chunk in chunks] | |
| try: | |
| scores = model.predict(pairs) | |
| except Exception as e: | |
| logger.error(f"Reranking failed: {e}", exc_info=True) | |
| # Fall back to original ordering | |
| for c in chunks: | |
| c["rerank_score"] = c.get("score", 0.0) | |
| return chunks[:top_k] | |
| # Attach rerank scores and sort | |
| for chunk, score in zip(chunks, scores): | |
| chunk["rerank_score"] = float(score) | |
| chunks.sort(key=lambda c: c["rerank_score"], reverse=True) | |
| reranked = chunks[:top_k] | |
| logger.debug( | |
| f"Reranked {len(chunks)} chunks → top {top_k} " | |
| f"(best score: {reranked[0]['rerank_score']:.4f})" | |
| ) | |
| return reranked | |