Spaces:
Running
Running
| """Author RAG Chatbot SaaS — Cross-Encoder Re-Ranker. | |
| Uses cross-encoder/ms-marco-MiniLM-L-6-v2 (free, local) to re-rank | |
| retrieved chunks by relevance to the original query. | |
| Significantly improves precision over cosine similarity alone. | |
| RULE: Keep top N chunks above minimum score threshold. | |
| BUG-6 fix: reranker.predict() is CPU-bound sync — wrapped in asyncio.to_thread(). | |
| """ | |
| import asyncio | |
| import structlog | |
| from app.config import get_settings | |
| from app.services.vector_store import RetrievedChunk | |
| logger = structlog.get_logger(__name__) | |
| cfg = get_settings() | |
| _reranker = None | |
| async def get_reranker(): | |
| """Lazily load and cache the cross-encoder re-ranker. | |
| Returns: | |
| Loaded CrossEncoder model. | |
| """ | |
| global _reranker | |
| if _reranker is None: | |
| from sentence_transformers import CrossEncoder | |
| logger.info("Loading cross-encoder re-ranker (first load)...") | |
| _reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2") | |
| logger.info("Cross-encoder re-ranker loaded successfully") | |
| return _reranker | |
| async def rerank_chunks( | |
| query: str, | |
| chunks: list[RetrievedChunk], | |
| top_n: int | None = None, | |
| min_score: float | None = None, | |
| ) -> list[RetrievedChunk]: | |
| """Re-rank retrieved chunks using cross-encoder scoring. | |
| Args: | |
| query: The original (non-rewritten) user query. | |
| chunks: List of RetrievedChunk from the retriever. | |
| top_n: Maximum chunks to keep after re-ranking. | |
| min_score: Minimum cross-encoder score to keep a chunk. | |
| Returns: | |
| Re-ranked and filtered list of chunks (best first). | |
| """ | |
| top_n = top_n or cfg.RAG_RERANK_TOP_N | |
| min_score = min_score or cfg.RAG_RERANK_MIN_SCORE | |
| if not chunks: | |
| return [] | |
| try: | |
| reranker = await get_reranker() | |
| # Build (query, chunk) pairs for cross-encoder | |
| pairs = [(query, chunk.text) for chunk in chunks] | |
| # BUG-6 fix: predict() is synchronous CPU-bound inference — offload to thread pool. | |
| scores = await asyncio.to_thread(reranker.predict, pairs) | |
| # Apply scores | |
| for chunk, score in zip(chunks, scores): | |
| chunk.rerank_score = float(score) | |
| # Sort by rerank score descending | |
| ranked = sorted(chunks, key=lambda c: c.rerank_score, reverse=True) | |
| # Filter by minimum score and limit to top_n | |
| filtered = [c for c in ranked if c.rerank_score >= min_score][:top_n] | |
| logger.debug( | |
| "Re-ranking complete", | |
| input_chunks=len(chunks), | |
| output_chunks=len(filtered), | |
| top_score=filtered[0].rerank_score if filtered else 0, | |
| ) | |
| return filtered | |
| except Exception as e: | |
| logger.error("Re-ranker failed, returning top-K by cosine score", error=str(e)) | |
| # Graceful fallback: return top chunks by initial similarity score | |
| return sorted(chunks, key=lambda c: c.score, reverse=True)[:top_n] | |