Spaces:
Sleeping
Sleeping
File size: 2,345 Bytes
8521f60 79bdbbe 8521f60 79bdbbe 8521f60 79bdbbe 8521f60 79bdbbe 3b8840f 79bdbbe 4dc151e 79bdbbe f868144 3b8840f f868144 79bdbbe 8521f60 79bdbbe 8521f60 79bdbbe 8521f60 79bdbbe 8521f60 79bdbbe 8521f60 |
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 |
from __future__ import annotations
import logging
from typing import List, Optional
from .base import Context, Retriever
from .bm25 import BM25Retriever
from .dense import DenseRetriever
logger = logging.getLogger(__name__)
class HybridRetriever(Retriever):
"""Combine BM25 and Dense retrievers by normalising and summing scores."""
def __init__(
self,
bm25_idx: str,
faiss_index: str,
doc_store: str,
*,
alpha: float = 0.5,
model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
embedder_cache: Optional[str] = None,
device: str = "cpu",
):
# 1) BM25 retriever
self.bm25 = BM25Retriever(bm25_idx, doc_store=doc_store)
# 2) Dense retriever
self.dense = DenseRetriever(
faiss_index=faiss_index,
doc_store=doc_store,
model_name=model_name,
embedder_cache=embedder_cache,
device=device,
)
if not 0 <= alpha <= 1:
raise ValueError("alpha must be in [0, 1]")
self.alpha = alpha
def retrieve(self, query: str, *, top_k: int = 5) -> List[Context]:
# 1) Get sparse hits
sparse_hits = self.bm25.retrieve(query, top_k=top_k)
sparse_dict = {ctx.id: ctx for ctx in sparse_hits}
# 2) Get dense hits
dense_hits = self.dense.retrieve(query, top_k=top_k)
dense_dict = {ctx.id: ctx for ctx in dense_hits}
# 3) Union of all IDs
all_ids = set(sparse_dict) | set(dense_dict)
merged: List[Context] = []
for doc_id in all_ids:
s_score = sparse_dict.get(doc_id, Context(doc_id, "", 0.0)).score
d_score = dense_dict.get(doc_id, Context(doc_id, "", 0.0)).score
combined_score = self.alpha * s_score + (1 - self.alpha) * d_score
# Prefer the text from whichever retriever has this doc_id present;
# if only one side has it, grab that text.
if doc_id in sparse_dict:
text = sparse_dict[doc_id].text
else:
text = dense_dict[doc_id].text
merged.append(Context(id=doc_id, text=text, score=combined_score))
# 4) Sort by score descending
merged.sort(key=lambda c: c.score, reverse=True)
return merged[:top_k]
|