Spaces:
Running on Zero
Running on Zero
| """Deterministic ranking and reciprocal-rank fusion.""" | |
| import json | |
| from pathlib import Path | |
| def rank_scores(page_ids: list[int], scores: list[float]) -> list[int]: | |
| if len(page_ids) != len(scores) or len(set(page_ids)) != len(page_ids): | |
| raise ValueError("Scores must align with unique page IDs.") | |
| return [ | |
| page for page, _ in sorted(zip(page_ids, scores, strict=True), key=lambda x: (-x[1], x[0])) | |
| ] | |
| def rrf(rankings: list[list[int]], k: int = 60) -> list[int]: | |
| """Fuse full rankings, with one-based ranks and page-ID tie breaking.""" | |
| if k < 1: | |
| raise ValueError("RRF k must be positive.") | |
| scores: dict[int, float] = {} | |
| for ranking in rankings: | |
| if len(set(ranking)) != len(ranking): | |
| raise ValueError("Each ranking must contain unique page IDs.") | |
| for rank, page in enumerate(ranking, start=1): | |
| scores[page] = scores.get(page, 0.0) + 1.0 / (k + rank) | |
| return rank_scores(list(scores), list(scores.values())) | |
| class PageRetriever: | |
| """Keep local data, indexes and the CPU encoder warm across queries.""" | |
| def __init__(self, subset: str = "hr", data_dir: Path = Path("data"), query_encoder=None): | |
| from visual_doc_rag.data import load_tables | |
| from visual_doc_rag.pages import load_pages | |
| demo_manifest = data_dir / subset / "demo.json" | |
| self.sample_rankings = {} | |
| self.query_encoder = query_encoder | |
| if demo_manifest.exists(): | |
| from datasets import Dataset | |
| payload = json.loads(demo_manifest.read_text(encoding="utf-8")) | |
| if payload["format_version"] != 1: | |
| raise ValueError("Unsupported demo bundle format.") | |
| self.source = payload["dataset"] | |
| self.sample_rankings = payload.get("sample_rankings", {}) | |
| corpus = json.loads((demo_manifest.parent / "corpus.json").read_text(encoding="utf-8")) | |
| self.tables = {"corpus": Dataset.from_list(corpus)} | |
| else: | |
| self.source, self.tables = load_tables(subset, data_dir) | |
| self.page_root, manifest = load_pages(subset, data_dir, self.source["revision"]) | |
| self.lookup = {page["corpus_id"]: page for page in manifest["pages"]} | |
| self.root = data_dir / subset / "indexes" | |
| self.data_dir = data_dir | |
| self.text = self.visual = self.encoder = None | |
| def search(self, query: str, method: str = "hybrid", top_k: int = 3) -> list[dict]: | |
| if method not in {"text", "visual", "hybrid"} or not query.strip() or top_k < 1: | |
| raise ValueError("Provide a nonempty query, a valid method and positive top-k.") | |
| sample = self.sample_rankings.get(" ".join(query.split())) | |
| if method == "hybrid" and sample and top_k <= len(sample): | |
| return self._results(sample[:top_k]) | |
| rankings = [] | |
| if method in {"text", "hybrid"}: | |
| from visual_doc_rag.text_index import TextIndex | |
| if self.text is None: | |
| self.text = TextIndex.load(self.root, self.source) | |
| rankings.append(self.text.search(query)) | |
| if method in {"visual", "hybrid"}: | |
| import torch | |
| from visual_doc_rag.visual_index import Encoder, VisualIndex | |
| torch.set_num_threads(4) | |
| if self.visual is None: | |
| self.visual = VisualIndex(self.root, self.source) | |
| if self.query_encoder is not None: | |
| embedding = self.query_encoder(query) | |
| else: | |
| if self.encoder is None: | |
| self.encoder = Encoder( | |
| self.data_dir, self.visual.manifest["signature"]["model"], "cpu" | |
| ) | |
| embedding = self.encoder.encode(query=query) | |
| rankings.append(self.visual.search(embedding)) | |
| ranking = rrf(rankings) if method == "hybrid" else rankings[0] | |
| return self._results(ranking[:top_k]) | |
| def _results(self, ranking): | |
| return [ | |
| { | |
| "rank": rank, | |
| **self.lookup[page], | |
| "image_path": str(self.page_root / self.lookup[page]["file"]), | |
| } | |
| for rank, page in enumerate(ranking, start=1) | |
| ] | |