Spaces:
Running on Zero
Running on Zero
| """BM25 over supplied page markdown; no benchmark answers or OCR pipeline.""" | |
| import json | |
| import re | |
| from pathlib import Path | |
| from time import perf_counter | |
| from rank_bm25 import BM25Okapi | |
| from visual_doc_rag.data import write_json | |
| from visual_doc_rag.retrieve import rank_scores | |
| def tokenize(text: str) -> list[str]: | |
| return re.findall(r"\w+", text.lower(), flags=re.UNICODE) | |
| class TextIndex: | |
| def __init__(self, page_ids: list[int], tokens: list[list[str]]): | |
| self.page_ids = page_ids | |
| self.tokens = tokens | |
| self.bm25 = BM25Okapi(tokens) # library defaults: k1=1.5, b=0.75, epsilon=0.25 | |
| def search(self, query: str) -> list[int]: | |
| return rank_scores(self.page_ids, self.bm25.get_scores(tokenize(query)).tolist()) | |
| def build(cls, corpus, root: Path, source: dict) -> dict: | |
| started = perf_counter() | |
| rows = corpus.remove_columns("image") | |
| payload = { | |
| "dataset": source, | |
| "page_ids": list(rows["corpus_id"]), | |
| "tokens": [tokenize(text or "") for text in rows["markdown"]], | |
| "tokenizer": "Unicode word regex, lowercase; no stemming/stopword removal", | |
| "parameters": {"k1": 1.5, "b": 0.75, "epsilon": 0.25}, | |
| } | |
| # Construct now as well, so build time includes the in-memory BM25 index. | |
| cls(payload["page_ids"], payload["tokens"]) | |
| write_json(root / "text.json", payload) | |
| result = { | |
| "build_seconds": perf_counter() - started, | |
| "index_bytes": (root / "text.json").stat().st_size, | |
| "pages": len(rows), | |
| "peak_vram_bytes": 0, | |
| } | |
| write_json(root / "text-build.json", result) | |
| return result | |
| def load(cls, root: Path, source: dict): | |
| payload = json.loads((root / "text.json").read_text(encoding="utf-8")) | |
| if payload["dataset"]["revision"] != source["revision"]: | |
| raise ValueError("Text index dataset revision differs from the downloaded data.") | |
| return cls(payload["page_ids"], payload["tokens"]) | |