Spaces:
Sleeping
Sleeping
File size: 1,821 Bytes
732b14f | 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 | """Reciprocal rank fusion for vector + BM25 retrieval lists."""
from __future__ import annotations
from rank_bm25 import BM25Okapi
from app.models.schemas import SearchResult
_RRF_K = 60
def tokenize(text: str) -> list[str]:
return [t for t in text.lower().split() if t]
def build_bm25_index(texts: list[str]) -> BM25Okapi | None:
corpus = [tokenize(t) for t in texts]
if not corpus:
return None
return BM25Okapi(corpus)
def bm25_search(
query: str,
*,
texts: list[str],
meta_rows: list[SearchResult],
k: int,
) -> list[SearchResult]:
"""Return top-k BM25 hits from ``meta_rows`` aligned with ``texts``."""
if not texts or not meta_rows:
return []
index = build_bm25_index(texts)
if index is None:
return []
scores = index.get_scores(tokenize(query))
ranked = sorted(
zip(meta_rows, scores, strict=True),
key=lambda pair: float(pair[1]),
reverse=True,
)
out: list[SearchResult] = []
for row, score in ranked[:k]:
out.append(row.model_copy(update={"score": float(score)}))
return out
def reciprocal_rank_fusion(
ranked_lists: list[list[SearchResult]],
*,
top_n: int,
rrf_k: int = _RRF_K,
) -> list[SearchResult]:
"""Merge multiple ranked lists with RRF; higher score = better."""
scores: dict[str, float] = {}
by_id: dict[str, SearchResult] = {}
for lst in ranked_lists:
for rank, item in enumerate(lst, start=1):
scores[item.chunk_id] = scores.get(item.chunk_id, 0.0) + 1.0 / (rrf_k + rank)
by_id[item.chunk_id] = item
ordered = sorted(scores.keys(), key=lambda cid: scores[cid], reverse=True)
return [
by_id[cid].model_copy(update={"score": scores[cid]})
for cid in ordered[:top_n]
]
|