Spaces:
Runtime error
Runtime error
| """Hybrid retrieval: vector search + BM25, fused with RRF, optionally re-ranked.""" | |
| import logging | |
| from . import store | |
| from .bm25 import bm25 | |
| from .config import settings | |
| log = logging.getLogger("rag.retrieval") | |
| def rrf(result_lists: list[list[dict]], k: int = 60) -> list[dict]: | |
| """Reciprocal Rank Fusion: merge several ranked lists into one.""" | |
| scores: dict[str, float] = {} | |
| items: dict[str, dict] = {} | |
| for lst in result_lists: | |
| for rank, hit in enumerate(lst): | |
| scores[hit["id"]] = scores.get(hit["id"], 0.0) + 1.0 / (k + rank + 1) | |
| items.setdefault(hit["id"], hit) | |
| fused = sorted(items.values(), key=lambda h: -scores[h["id"]]) | |
| return [{**h, "score": scores[h["id"]]} for h in fused] | |
| def _interleave(groups: list[list[dict]], limit: int) -> list[dict]: | |
| """Round-robin across documents so one big PDF cannot crowd out the others.""" | |
| out: list[dict] = [] | |
| i = 0 | |
| while len(out) < limit and any(i < len(g) for g in groups): | |
| for g in groups: | |
| if i < len(g) and len(out) < limit: | |
| out.append(g[i]) | |
| i += 1 | |
| return out | |
| def hybrid_search(query: str, doc_ids: list[str], pages: list[int] | None = None, top_k: int | None = None) -> list[dict]: | |
| top_k = top_k or settings.top_k | |
| vec = store.embed([query])[0] | |
| if 1 < len(doc_ids) <= settings.max_doc_groups: | |
| groups = [[d] for d in doc_ids] # search each document on its own | |
| per_group = settings.per_doc_k | |
| else: | |
| groups = [list(doc_ids)] | |
| per_group = settings.candidate_k | |
| ranked_groups: list[list[dict]] = [] | |
| for g in groups: | |
| dense = store.search(vec, store.make_filter(g, pages), per_group) | |
| sparse = bm25.search(query, g, pages, per_group) | |
| dense_score = {h["id"]: h["score"] for h in dense} | |
| fused = rrf([dense, sparse])[:per_group] | |
| # keep the raw cosine similarity: it is comparable across documents, RRF scores are not | |
| ranked_groups.append([{**h, "dense": dense_score.get(h["id"], 0.0)} for h in fused]) | |
| candidates = [h for grp in ranked_groups for h in grp] | |
| if not candidates: | |
| return [] | |
| reranker = store.get_reranker() | |
| if reranker is not None and len(candidates) > 1: | |
| scores = reranker.predict([(query, c["text"]) for c in candidates], batch_size=16, show_progress_bar=False) | |
| ranked = sorted(zip(candidates, scores), key=lambda x: -float(x[1])) | |
| return [{**c, "score": float(s)} for c, s in ranked[:top_k]] | |
| # no re-ranker: put the documents with the best semantic match first, then round-robin | |
| ranked_groups = [g for g in ranked_groups if g] | |
| ranked_groups.sort(key=lambda g: -max(h["dense"] for h in g)) | |
| return _interleave(ranked_groups, top_k) | |
| def fetch_pages(doc_ids: list[str], pages: list[int]) -> list[dict]: | |
| """Every chunk on the given pages (exact lookup, no similarity search).""" | |
| chunks = store.scroll_all(store.make_filter(doc_ids, pages)) | |
| chunks.sort(key=lambda c: (c["filename"], c["page"], c.get("chunk_index", 0))) | |
| return chunks | |