src/visual_doc_rag/data.py

"""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