rpx_benchmark.dataset_hub.croissant

  1"""Stage the Croissant metadata JSON into the upload tree.
  2
  3Croissant (https://mlcommons.org/working-groups/croissant/) is the ML
  4metadata standard HuggingFace renders alongside dataset cards and that
  5search platforms (Kaggle, Papers-with-Code) parse. We already authored
  6``rpx_benchmark/dataset_hub/croissant_template.json`` against
  7an earlier dataset URL; this helper:
  8
  91. **Copies** the source JSON into ``<staging>/rpx_croissant.json``.
 102. **Patches** the `url`, `version`, and `citeAs` fields to point at the
 11   HF repo (the source pre-dates the HF repo decision).
 123. **Annotates** the `description` so consumers know that per-frame
 13   records now live inside per-modality tar shards rather than as loose
 14   files (the field-level path descriptions in the JSON still describe
 15   the logical type, but the on-disk shape changed under our packer).
 16
 17A full path-by-path schema reconciliation against the new tar layout is
 18left to the dataset owner — it's a deeper edit that depends on what
 19loaders downstream consumers actually use, and is not blocking for the
 20v1 upload.
 21"""
 22
 23from __future__ import annotations
 24
 25import json
 26from dataclasses import dataclass
 27from pathlib import Path
 28from typing import Optional
 29
 30from ..exceptions import ConfigError, DatasetError
 31from ..logging_utils import get_logger
 32from .recipes import DEFAULT_REPO_ID
 33
 34log = get_logger(__name__)
 35
 36
 37# Default source of the Croissant template. The JSON ships inside the
 38# package so pip-installed users get it without needing any local
 39# checkout outside the package.
 40def _default_croissant_src() -> Path:
 41    """``<package>/dataset_hub/croissant_template.json``."""
 42    return Path(__file__).resolve().parent / "croissant_template.json"
 43
 44
 45@dataclass(frozen=True)
 46class CroissantPatch:
 47    """Knobs for the staging step."""
 48
 49    repo_id: str = DEFAULT_REPO_ID
 50    version: str = "1.0.0"
 51    cite_as: Optional[str] = None
 52    annotate_layout: bool = True
 53
 54
 55def _hf_url(repo_id: str) -> str:
 56    return f"https://huggingface.co/datasets/{repo_id}"
 57
 58
 59_LAYOUT_NOTE = (
 60    "\n\nNote on on-disk shape: per-frame records are physically packed "
 61    "into per-(scene, phase, modality) tar shards on the HF repo "
 62    "(scenes/<scene>/<phase>/<modality>.tar for raw modalities; "
 63    "labels/<modality>/v<N>.tar for versioned labels). The "
 64    "rpx_benchmark.dataset_hub downloader resolves selective task+split "
 65    "downloads from these shards transparently."
 66)
 67
 68
 69def stage_croissant(
 70    staging_root: Path,
 71    src: Optional[Path] = None,
 72    patch: Optional[CroissantPatch] = None,
 73    overwrite: bool = False,
 74) -> Path:
 75    """Copy + patch the Croissant JSON into ``<staging>/rpx_croissant.json``.
 76
 77    Returns the destination path.
 78    """
 79    src = Path(src) if src else _default_croissant_src()
 80    patch = patch or CroissantPatch()
 81    if not src.is_file():
 82        raise ConfigError(
 83            f"croissant source JSON does not exist: {src}",
 84            hint=(
 85                "Pass src= explicitly, or check that "
 86                "croissant_template.json is present in the package "
 87                "in the repo."
 88            ),
 89        )
 90
 91    out_path = Path(staging_root) / "rpx_croissant.json"
 92    if out_path.exists() and not overwrite:
 93        raise DatasetError(
 94            f"refusing to overwrite existing croissant: {out_path}",
 95            hint="Pass overwrite=True or remove the file first.",
 96        )
 97
 98    payload = json.loads(src.read_text(encoding="utf-8"))
 99
100    payload["url"] = _hf_url(patch.repo_id)
101    payload["version"] = patch.version
102    if patch.cite_as is not None:
103        payload["citeAs"] = patch.cite_as
104
105    if patch.annotate_layout:
106        desc = payload.get("description", "")
107        if _LAYOUT_NOTE.strip() not in desc:
108            payload["description"] = desc + _LAYOUT_NOTE
109
110    out_path.parent.mkdir(parents=True, exist_ok=True)
111    out_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
112    log.info("staged croissant → %s (%d bytes)", out_path, out_path.stat().st_size)
113    return out_path