File size: 2,610 Bytes
b7f3196 |
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 |
from typing import Dict, List
from retriever import Retriever
from retriever.rrf import rrf
from team.interfaces import Candidate
from pathlib import Path
def _default_corpora_config() -> Dict[str, dict]:
return {
"medical_qa": {"path":"data/corpora/medical_qa.jsonl",
"text_fields":["question","answer","title"]},
"miriad": {"path":"data/corpora/miriad_text.jsonl",
"text_fields":["question","answer","title"]},
"pubmed": {"path":"data/corpora/pubmed.jsonl",
"text_fields":["contents","title"]},
"unidoc": {"path":"data/corpora/unidoc_qa.jsonl",
"text_fields":["question","answer","title"]},
}
def _available(cfg: Dict[str, dict]) -> Dict[str, dict]:
return {k:v for k,v in cfg.items() if Path(v["path"]).exists()}
def get_candidates(
query: str,
retriever: Retriever,
k_retrieve: int = 50,
) -> List[Candidate]:
"""
Returns top-N fused candidates with component scores (bm25, dense, rrf).
"""
r = retriever
# get separate result lists (doc, score)
bm = r.bm25.search(query, k=max(k_retrieve, 100))
de = r.dense.search(query, k=max(k_retrieve, 100))
# maps for score lookup
bm_map = {d.id: float(s) for d, s in bm}
de_map = {d.id: float(s) for d, s in de}
# fuse and pick candidate set
fused = rrf([bm, de], k=max(k_retrieve, 50))
# compute RRF per candidate using rank positions
K = 60
bm_rank = {d.id:i for i,(d,_) in enumerate(bm)}
de_rank = {d.id:i for i,(d,_) in enumerate(de)}
out: List[Candidate] = []
for doc, _ in fused[:k_retrieve]:
rrf_score = 0.0
if doc.id in bm_rank:
rrf_score += 1.0 / (K + bm_rank[doc.id] + 1)
if doc.id in de_rank:
rrf_score += 1.0 / (K + de_rank[doc.id] + 1)
out.append(Candidate(
id=doc.id,
title=doc.title or "",
text=doc.text,
meta=doc.meta or {},
bm25=bm_map.get(doc.id, 0.0),
dense=de_map.get(doc.id, 0.0),
rrf=rrf_score,
))
# baseline order: RRF
out.sort(key=lambda c: c.rrf, reverse=True)
return out
#how to call/run below for everyone
# from team.candidates import get_candidates
# q = "worst headache of my life with fever and stiff neck"
# cands = get_candidates(q, k_retrieve=60) # returns List[Candidate]
# for c in cands[:3]:
# print(c.id, c.bm25, c.dense, c.rrf, c.title)
|