Spaces:
Sleeping
Sleeping
v2: cited answers with hybrid retrieval + cross-encoder rerank, Qwen3-1.7B on CPU, retrieval trace
a072368 verified | """Hybrid retrieval: dense + BM25, fused with Reciprocal Rank Fusion, then a cross-encoder rerank.""" | |
| from dataclasses import dataclass, field | |
| from typing import Protocol | |
| from .chunking import Chunk | |
| from .index import Index | |
| class Reranker(Protocol): | |
| def score(self, query: str, passages: list[str]) -> list[float]: ... | |
| class Hit: | |
| chunk: Chunk | |
| vector_rank: int | None = None | |
| vector_score: float | None = None | |
| keyword_rank: int | None = None | |
| keyword_score: float | None = None | |
| fused_rank: int = 0 | |
| fused_score: float = 0.0 | |
| rerank_score: float | None = None | |
| rank: int = 0 # final position (1-based) | |
| class Retrieval: | |
| query: str | |
| hits: list[Hit] = field(default_factory=list) # final order, top_k | |
| candidates: list[Hit] = field(default_factory=list) # fused order, before reranking | |
| def confidence(self) -> float | None: | |
| return self.hits[0].rerank_score if self.hits else None | |
| def rrf(rankings: list[list[str]], k: int = 60) -> list[tuple[str, float]]: | |
| """Reciprocal Rank Fusion: sum of 1 / (k + rank) over every list a document appears in.""" | |
| scores: dict[str, float] = {} | |
| for ranking in rankings: | |
| for pos, cid in enumerate(ranking, start=1): | |
| scores[cid] = scores.get(cid, 0.0) + 1.0 / (k + pos) | |
| return sorted(scores.items(), key=lambda kv: -kv[1]) | |
| class Retriever: | |
| def __init__( | |
| self, | |
| index: Index, | |
| reranker: Reranker | None = None, | |
| candidates: int = 20, | |
| rerank_pool: int = 12, | |
| top_k: int = 5, | |
| rrf_k: int = 60, | |
| mode: str = "hybrid", # "hybrid" | "vector" | "keyword" | |
| ) -> None: | |
| self.index, self.reranker = index, reranker | |
| self.candidates, self.rerank_pool, self.top_k, self.rrf_k, self.mode = ( | |
| candidates, | |
| rerank_pool, | |
| top_k, | |
| rrf_k, | |
| mode, | |
| ) | |
| def search(self, query: str) -> Retrieval: | |
| vec = self.index.vector_search(query, self.candidates) if self.mode in ("hybrid", "vector") else [] | |
| kw = self.index.keyword_search(query, self.candidates) if self.mode in ("hybrid", "keyword") else [] | |
| hits: dict[str, Hit] = {} | |
| for pos, (cid, s) in enumerate(vec, start=1): | |
| hits.setdefault(cid, Hit(self.index.chunks[cid])).__dict__.update(vector_rank=pos, vector_score=s) | |
| for pos, (cid, s) in enumerate(kw, start=1): | |
| hits.setdefault(cid, Hit(self.index.chunks[cid])).__dict__.update(keyword_rank=pos, keyword_score=s) | |
| fused = rrf([[c for c, _ in vec], [c for c, _ in kw]], self.rrf_k) | |
| candidates = [] | |
| for pos, (cid, s) in enumerate(fused, start=1): | |
| h = hits[cid] | |
| h.fused_rank, h.fused_score = pos, s | |
| candidates.append(h) | |
| final = candidates[: self.rerank_pool] | |
| if self.reranker and final: | |
| scores = self.reranker.score(query, [h.chunk.text for h in final]) | |
| for h, s in zip(final, scores): | |
| h.rerank_score = float(s) | |
| final = sorted(final, key=lambda h: -h.rerank_score) | |
| final = final[: self.top_k] | |
| for pos, h in enumerate(final, start=1): | |
| h.rank = pos | |
| return Retrieval(query=query, hits=final, candidates=candidates[: self.rerank_pool]) | |