Spaces:
Running on Zero
Running on Zero
| """Download a pinned snapshot and inspect the real Parquet tables locally.""" | |
| import json | |
| from collections import Counter | |
| from datetime import UTC, datetime | |
| from pathlib import Path | |
| from datasets import Image, load_dataset | |
| from huggingface_hub import HfApi, snapshot_download | |
| SUBSETS = {"hr": "vidore/vidore_v3_hr", "industrial": "vidore/vidore_v3_industrial"} | |
| CONFIGS = ("corpus", "queries", "qrels", "documents_metadata") | |
| def write_json(path: Path, value: dict) -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| temporary = path.with_suffix(".tmp") | |
| temporary.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") | |
| temporary.replace(path) | |
| def download(subset: str, data_dir: Path, revision: str = "main") -> dict: | |
| """Resolve the revision once, download only tables/card, then publish a manifest.""" | |
| repo_id = SUBSETS[subset] | |
| root = data_dir / subset | |
| info = HfApi().dataset_info(repo_id, revision=revision) | |
| sha = info.sha | |
| if not sha: | |
| raise ValueError("Hugging Face did not return a dataset revision.") | |
| snapshot = root / "snapshots" / sha | |
| snapshot_download( | |
| repo_id=repo_id, | |
| repo_type="dataset", | |
| revision=sha, | |
| local_dir=snapshot, | |
| cache_dir=root / "cache" / "hub", | |
| allow_patterns=["README.md", *[f"{config}/*.parquet" for config in CONFIGS]], | |
| max_workers=2, | |
| ) | |
| for config in CONFIGS: | |
| if not list((snapshot / config).glob("*.parquet")): | |
| raise ValueError(f"No Parquet files found for {config}.") | |
| manifest = { | |
| "repo_id": repo_id, | |
| "revision": sha, | |
| "downloaded_at": datetime.now(UTC).isoformat(), | |
| "snapshot": str(snapshot.relative_to(root)), | |
| "configs": list(CONFIGS), | |
| "parquet_bytes": sum(p.stat().st_size for p in snapshot.rglob("*.parquet")), | |
| } | |
| write_json(root / "manifest.json", manifest) | |
| return manifest | |
| def load_tables(subset: str, data_dir: Path) -> tuple[dict, dict]: | |
| root = data_dir / subset | |
| manifest_path = root / "manifest.json" | |
| if not manifest_path.exists(): | |
| raise ValueError(f"Missing {manifest_path}. Run `vdr download --subset {subset}` first.") | |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) | |
| snapshot = root / manifest["snapshot"] | |
| tables = {} | |
| for config in CONFIGS: | |
| files = sorted((snapshot / config).glob("*.parquet")) | |
| if not files: | |
| raise ValueError(f"Missing {config} Parquet files. Run download again.") | |
| tables[config] = load_dataset( | |
| "parquet", | |
| data_files={"test": [str(p) for p in files]}, | |
| split="test", | |
| cache_dir=str(root / "cache" / "datasets" / manifest["revision"]), | |
| ) | |
| return manifest, tables | |
| def inspect_tables(tables: dict) -> dict: | |
| """Summarize metadata and validate joins without exporting query/answer text.""" | |
| corpus = tables["corpus"].remove_columns("image") | |
| queries, qrels, documents = (tables[name] for name in CONFIGS[1:]) | |
| page_ids = list(corpus["corpus_id"]) | |
| query_ids = list(queries["query_id"]) | |
| page_set, query_set = set(page_ids), set(query_ids) | |
| english_ids = { | |
| row["query_id"] for row in queries if row["language"].strip().lower() in {"en", "english"} | |
| } | |
| positive_queries = {row["query_id"] for row in qrels if row["score"] > 0} | |
| english_qrels = [row for row in qrels if row["query_id"] in english_ids] | |
| pairs = [(row["query_id"], row["corpus_id"]) for row in qrels] | |
| checks = { | |
| "unique_corpus_ids": len(page_ids) == len(page_set), | |
| "unique_query_ids": len(query_ids) == len(query_set), | |
| "unique_qrel_pairs": len(pairs) == len(set(pairs)), | |
| "qrels_reference_existing_pages": all(row["corpus_id"] in page_set for row in qrels), | |
| "qrels_reference_existing_queries": all(row["query_id"] in query_set for row in qrels), | |
| "corpus_documents_have_metadata": set(corpus["doc_id"]) <= set(documents["doc_id"]), | |
| "english_queries_present": bool(english_ids), | |
| "english_queries_have_positive_qrels": english_ids <= positive_queries, | |
| "nonempty_english_queries": all( | |
| row["query"].strip() for row in queries if row["query_id"] in english_ids | |
| ), | |
| } | |
| return { | |
| "rows": {name: len(table) for name, table in tables.items()}, | |
| "features": {name: table.features.to_dict() for name, table in tables.items()}, | |
| "query_languages": dict(Counter(queries["language"])), | |
| "english_queries": len(english_ids), | |
| "english_qrels": len(english_qrels), | |
| "english_positive_pages": len( | |
| {row["corpus_id"] for row in english_qrels if row["score"] > 0} | |
| ), | |
| "qrel_scores": dict(Counter(qrels["score"])), | |
| "english_qrel_content_types": dict( | |
| Counter(kind for row in english_qrels for kind in row["content_type"]) | |
| ), | |
| "document_visual_types": dict( | |
| Counter(kind for row in documents for kind in row["visual_types"]) | |
| ), | |
| "document_licenses": dict(Counter(documents["license"])), | |
| "pages_with_nonempty_markdown": sum( | |
| bool(text and text.strip()) for text in corpus["markdown"] | |
| ), | |
| "markdown_characters": sum(len(text or "") for text in corpus["markdown"]), | |
| "page_number_range": [min(corpus["page_number_in_doc"]), max(corpus["page_number_in_doc"])], | |
| "integrity_checks": checks, | |
| } | |
| def save_samples(tables: dict, root: Path) -> list[dict]: | |
| """Decode a few pages selected from visual evidence labels for manual review.""" | |
| from PIL import Image as PILImage | |
| from PIL import ImageDraw | |
| visual_ids = [] | |
| for row in tables["qrels"]: | |
| types = " ".join(row["content_type"]).lower() | |
| if ( | |
| any(kind in types for kind in ("table", "chart", "figure", "image")) | |
| and row["corpus_id"] not in visual_ids | |
| ): | |
| visual_ids.append(row["corpus_id"]) | |
| if len(visual_ids) == 6: | |
| break | |
| corpus = tables["corpus"].cast_column("image", Image(decode=True)) | |
| ids = list(corpus["corpus_id"]) | |
| selected = visual_ids or ids[:6] | |
| output = root / "samples" | |
| output.mkdir(parents=True, exist_ok=True) | |
| sheet = PILImage.new("RGB", (960, 1440), "#dddddd") | |
| draw = ImageDraw.Draw(sheet) | |
| samples = [] | |
| for index, page_id in enumerate(selected): | |
| row = corpus[ids.index(page_id)] | |
| with row["image"] as image: | |
| sample = image.convert("RGB") | |
| original_size = list(sample.size) | |
| sample.thumbnail((900, 1200)) | |
| sample.save(output / f"{page_id}.jpg", quality=85) | |
| sample.thumbnail((460, 440)) | |
| x, y = (index % 2) * 480, (index // 2) * 480 | |
| sheet.paste(sample, (x, y + 24)) | |
| draw.text((x + 8, y + 6), f"corpus_id={page_id}", fill="black") | |
| samples.append({"corpus_id": page_id, "image_size": original_size}) | |
| sheet.save(output / "contact-sheet.jpg", quality=85) | |
| return samples | |
| def inspect(subset: str, data_dir: Path) -> dict: | |
| manifest, tables = load_tables(subset, data_dir) | |
| report = inspect_tables(tables) | |
| report["dataset"] = manifest | |
| report["samples"] = save_samples(tables, data_dir / subset) | |
| write_json(data_dir / subset / "inspection.json", report) | |
| return report | |