Spaces:
Running
Running
| """Hybrid search combining Dense Vector Search and BM25 Keyword Search. | |
| Uses Reciprocal Rank Fusion (RRF) to merge the ranked result lists from both | |
| retrieval methods into a single, high-quality candidate list. | |
| Why Hybrid Search? | |
| - Dense vectors understand semantic similarity (e.g. 'space telescope' -> 'astronomy'). | |
| - BM25 understands exact identifiers and keywords (e.g. 'NIRCam', 'L2', '15 millikelvin'). | |
| - Together, they significantly outperform either method on its own. | |
| """ | |
| from collections import defaultdict | |
| def reciprocal_rank_fusion(dense_results, bm25_results, dense_weight=0.6, bm25_weight=0.4, k=60): | |
| """Merge two ranked lists using Reciprocal Rank Fusion (RRF). | |
| Formula for each document: | |
| Score = sum( weight / (k + rank) ) | |
| Args: | |
| dense_results: List of dicts from vector search [{"chunk": ..., "score": ...}]. | |
| bm25_results: List of dicts from BM25 search. | |
| dense_weight: Relative importance of vector search (default: 0.6). | |
| bm25_weight: Relative importance of keyword search (default: 0.4). | |
| k: Smoothing constant to dampen extreme high ranks (standard is 60). | |
| Returns: | |
| List of fused dicts sorted by combined RRF score. | |
| """ | |
| rrf_scores = defaultdict(float) | |
| chunks_by_id = {} | |
| dense_ranks = {} | |
| bm25_ranks = {} | |
| # Accumulate dense ranks | |
| for rank, item in enumerate(dense_results, start=1): | |
| chunk = item["chunk"] | |
| cid = chunk.chunk_id | |
| chunks_by_id[cid] = chunk | |
| dense_ranks[cid] = rank | |
| rrf_scores[cid] += dense_weight * (1.0 / (k + rank)) | |
| # Accumulate BM25 ranks | |
| for rank, item in enumerate(bm25_results, start=1): | |
| chunk = item["chunk"] | |
| cid = chunk.chunk_id | |
| chunks_by_id[cid] = chunk | |
| bm25_ranks[cid] = rank | |
| rrf_scores[cid] += bm25_weight * (1.0 / (k + rank)) | |
| # Sort descending by fused score | |
| sorted_items = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True) | |
| fused_results = [] | |
| for cid, combined_score in sorted_items: | |
| fused_results.append({ | |
| "chunk": chunks_by_id[cid], | |
| "score": round(combined_score, 5), | |
| "dense_rank": dense_ranks.get(cid), | |
| "bm25_rank": bm25_ranks.get(cid), | |
| "source": "hybrid", | |
| }) | |
| return fused_results | |
| class HybridRetriever: | |
| """Orchestrates dense vector search and BM25 search.""" | |
| def __init__(self, vector_store, bm25_index, dense_weight=0.6, bm25_weight=0.4): | |
| self.vector_store = vector_store | |
| self.bm25_index = bm25_index | |
| self.dense_weight = dense_weight | |
| self.bm25_weight = bm25_weight | |
| def retrieve(self, query, top_k=4, dense_k=10, bm25_k=10): | |
| """Retrieve the best matching chunks using hybrid search. | |
| Returns: | |
| Dict containing: | |
| - 'results': top_k fused results | |
| - 'dense_results': raw vector results | |
| - 'bm25_results': raw BM25 results | |
| """ | |
| dense_hits = self.vector_store.search(query, top_k=dense_k) | |
| bm25_hits = self.bm25_index.search(query, top_k=bm25_k) | |
| fused = reciprocal_rank_fusion( | |
| dense_results=dense_hits, | |
| bm25_results=bm25_hits, | |
| dense_weight=self.dense_weight, | |
| bm25_weight=self.bm25_weight, | |
| ) | |
| return { | |
| "results": fused[:top_k], | |
| "dense_results": dense_hits, | |
| "bm25_results": bm25_hits, | |
| } | |