rpx_benchmark.data.load

  1"""One-line ``datasets.load_dataset`` entry point for RPX.
  2
  3Example
  4-------
  5
  6::
  7
  8    from rpx_benchmark.data import load_hf
  9
 10    ds = load_hf("monocular_depth", split="hard")
 11    for batch in ds:
 12        # each sample is a full rpx_benchmark.Sample
 13        ...
 14"""
 15
 16from __future__ import annotations
 17
 18from typing import Any
 19
 20from ..api import Difficulty, TaskType
 21from ..exceptions import DownloadError
 22from ..hub import DEFAULT_REPO_ID, DEFAULT_REVISION
 23from ..logging_utils import get_logger
 24from .hf_bridge import RPXHFBridge
 25
 26log = get_logger(__name__)
 27
 28def _load_dataset(*args: Any, **kwargs: Any) -> Any:
 29    """Lazy import of :func:`datasets.load_dataset`.
 30
 31    Kept lazy so ``import rpx_benchmark.data.load`` does not force a
 32    pyarrow / datasets import at package load time for users who only
 33    wanted the convenience re-export surface.
 34    """
 35    try:
 36        from datasets import load_dataset  # noqa: PLC0415
 37    except ImportError as e:  # pragma: no cover
 38        raise ImportError(
 39            "rpx_benchmark.data.load_hf requires the `datasets` library. "
 40            "Install with: pip install 'rpx-benchmark[hf-datasets]'"
 41        ) from e
 42    return load_dataset(*args, **kwargs)
 43
 44
 45def load_hf(
 46    task: TaskType | str,
 47    split: Difficulty | str = Difficulty.HARD,
 48    *,
 49    repo_id: str = DEFAULT_REPO_ID,
 50    revision: str | None = DEFAULT_REVISION,
 51    cache_dir: str | None = None,
 52    streaming: bool = False,
 53    batch_size: int = 1,
 54    **load_dataset_kwargs: Any,
 55) -> RPXHFBridge:
 56    """Load an RPX task slice as an :class:`RPXHFBridge`.
 57
 58    Parameters
 59    ----------
 60    task : TaskType or str
 61        Task config name on the HF repo. Uses the enum ``.value`` if a
 62        :class:`TaskType` is passed.
 63    split : Difficulty or str, default ``"hard"``
 64        ESD split to fetch. Maps to ``datasets.load_dataset``'s
 65        ``split`` kwarg.
 66    repo_id : str
 67        Hugging Face dataset repo id. Defaults to
 68        ``IRVLUTD/RPX`` (overridable via ``RPX_HF_REPO``).
 69    revision : str, optional
 70        Git revision (tag, branch, or commit) to pin. Defaults to the
 71        consolidated audited release; override it with ``RPX_HF_REVISION``.
 72    cache_dir : str, optional
 73        Hugging Face cache directory. When None the library default is
 74        used (``~/.cache/huggingface``).
 75    streaming : bool, default False
 76        Pass through to ``datasets.load_dataset``. Enables row-level
 77        streaming from the Hub without materialising the full Arrow
 78        tables locally.
 79    batch_size : int, default 1
 80        Forwarded into the returned :class:`RPXHFBridge`.
 81    **load_dataset_kwargs
 82        Any extra kwargs are forwarded to ``datasets.load_dataset``.
 83
 84    Returns
 85    -------
 86    RPXHFBridge
 87        An iterable wrapper that yields ``list[Sample]`` batches.
 88
 89    Raises
 90    ------
 91    DownloadError
 92        If the underlying ``datasets.load_dataset`` call fails.
 93    """
 94    task_enum = TaskType(task) if isinstance(task, str) else task
 95    split_name = split.value if isinstance(split, Difficulty) else str(split)
 96
 97    log.info(
 98        "loading HF dataset repo=%s config=%s split=%s streaming=%s",
 99        repo_id,
100        task_enum.value,
101        split_name,
102        streaming,
103    )
104    try:
105        hf_ds = _load_dataset(
106            repo_id,
107            name=task_enum.value,
108            split=split_name,
109            revision=revision,
110            cache_dir=cache_dir,
111            streaming=streaming,
112            **load_dataset_kwargs,
113        )
114    except Exception as e:
115        raise DownloadError(
116            f"datasets.load_dataset failed for {repo_id}/{task_enum.value} "
117            f"(split={split_name}): {e}",
118            hint=(
119                "Check the repo id and your HF_TOKEN for private datasets. "
120                "For offline reuse set HF_DATASETS_OFFLINE=1 after a first "
121                "online load populates the cache."
122            ),
123        ) from e
124
125    return RPXHFBridge(hf_dataset=hf_ds, task=task_enum, batch_size=batch_size)