Spaces:
Running on Zero
Running on Zero
File size: 4,231 Bytes
86e797b | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """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)
]
|