Spaces:
Running
Running
| """ | |
| reranker.py | |
| ----------- | |
| Re-ranks retrieved chunks using a cross-encoder model for better relevance ordering. | |
| Cross-encoders (unlike bi-encoders) compute similarity by processing the query | |
| and document together, providing more accurate ranking than vector similarity alone. | |
| Uses: cross-encoder/ms-marco-MiniLM-L-12-v2 (fast, accurate, lightweight) | |
| """ | |
| import logging | |
| from typing import List, Tuple | |
| from langchain.schema import Document | |
| from sentence_transformers import CrossEncoder | |
| from app.config import RERANKER_MODEL_NAME, RERANKER_TOP_K | |
| logger = logging.getLogger(__name__) | |
| class Reranker: | |
| """ | |
| Re-ranks retrieved document chunks using a cross-encoder model. | |
| Args: | |
| model_name: HuggingFace model ID for the cross-encoder (default from config). | |
| top_k: Number of results to return after re-ranking. | |
| """ | |
| def __init__( | |
| self, | |
| model_name: str = RERANKER_MODEL_NAME, | |
| top_k: int = RERANKER_TOP_K, | |
| ) -> None: | |
| self.model_name = model_name | |
| self.top_k = top_k | |
| self._model: CrossEncoder | None = None | |
| def _load_model(self) -> CrossEncoder: | |
| """ | |
| Lazily load the cross-encoder model on first use. | |
| """ | |
| if self._model is None: | |
| logger.info(f"Loading cross-encoder model: {self.model_name}") | |
| self._model = CrossEncoder(self.model_name) | |
| return self._model | |
| def rerank( | |
| self, | |
| query: str, | |
| documents: List[Tuple[Document, float]], | |
| top_k: int | None = None, | |
| ) -> List[Tuple[Document, float]]: | |
| """ | |
| Re-rank retrieved documents by relevance to the query. | |
| Uses cross-encoder scores instead of raw vector similarity. | |
| Lower cross-encoder scores are better (0=irrelevant, 1=highly relevant in most cases, | |
| but the scale depends on the model). | |
| Args: | |
| query: User's natural-language question. | |
| documents: List of (Document, vector_score) tuples from initial retrieval. | |
| top_k: Number of results to return (falls back to self.top_k). | |
| Returns: | |
| List of (Document, cross_encoder_score) tuples, sorted by relevance. | |
| """ | |
| if not documents: | |
| logger.debug("No documents to rerank.") | |
| return [] | |
| k = top_k or self.top_k | |
| model = self._load_model() | |
| # Extract document texts and original scores | |
| doc_list = [doc for doc, _ in documents] | |
| original_scores = {doc.page_content: score for doc, score in documents} | |
| # Prepare pairs for cross-encoder: (query, document_text) | |
| pairs = [[query, doc.page_content] for doc in doc_list] | |
| try: | |
| # Get cross-encoder scores | |
| # Returns a list of scores; higher is better for most models | |
| ce_scores = model.predict(pairs) | |
| # Combine documents with their cross-encoder scores | |
| ranked = list(zip(doc_list, ce_scores)) | |
| # Sort by cross-encoder score (descending) | |
| ranked.sort(key=lambda x: x[1], reverse=True) | |
| # Return top-k | |
| result = ranked[:k] | |
| logger.debug( | |
| f"Reranked {len(documents)} documents → top {len(result)} by cross-encoder" | |
| ) | |
| return result | |
| except Exception as exc: | |
| logger.error(f"Cross-encoder reranking failed: {exc}") | |
| # Fall back to original vector scores if reranking fails | |
| return documents[:k] | |
| def rerank_and_get_documents( | |
| self, | |
| query: str, | |
| documents: List[Tuple[Document, float]], | |
| top_k: int | None = None, | |
| ) -> List[Document]: | |
| """ | |
| Convenience wrapper — returns only Document objects (no scores). | |
| """ | |
| return [doc for doc, _ in self.rerank(query, documents, top_k=top_k)] | |