Spaces:
Running
Running
File size: 3,509 Bytes
e43c82d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """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,
}
|