rpx_benchmark.data.torch_dataloader

 1"""PyTorch ``DataLoader`` adapter for RPX datasets.
 2
 3Wraps either :class:`rpx_benchmark.loader.RPXDataset` or
 4:class:`rpx_benchmark.data.hf_bridge.RPXHFBridge` so users can plug
 5into the standard PyTorch training loop while still receiving
 6``list[Sample]`` batches (which adapters and the runner expect).
 7
 8The collate function is deliberately trivial — it returns the list of
 9samples unchanged. Task-specific tensor stacking belongs in the
10:class:`InputAdapter`, not here, so the dataloader stays model-agnostic.
11"""
12
13from __future__ import annotations
14
15from typing import Any, Iterable, List, Sequence
16
17from ..api import Sample
18
19
20def rpx_collate(batch: Sequence[Any]) -> List[Sample]:
21    """Collate a batch of either ``Sample`` objects or already-batched lists.
22
23    ``RPXDataset`` / ``RPXHFBridge`` produce ``list[Sample]`` per
24    iteration, whereas PyTorch's ``DataLoader`` expects a dataset whose
25    ``__getitem__`` returns one item. We therefore flatten in case the
26    inner dataset is already batched.
27    """
28    flat: List[Sample] = []
29    for item in batch:
30        if isinstance(item, Sample):
31            flat.append(item)
32        elif isinstance(item, Sequence):
33            flat.extend(s for s in item if isinstance(s, Sample))
34    return flat
35
36
37def to_torch_dataloader(
38    rpx_iterable: Iterable[Any],
39    *,
40    num_workers: int = 0,
41    pin_memory: bool = False,
42) -> Any:
43    """Wrap an RPX iterable in a ``torch.utils.data.DataLoader``.
44
45    The returned DataLoader yields ``list[Sample]`` batches of the same
46    size as the wrapped iterable's ``batch_size``.
47
48    Parameters
49    ----------
50    rpx_iterable : Iterable
51        Any object with ``__iter__`` yielding ``list[Sample]``.
52    num_workers : int, default 0
53        Passed to ``DataLoader``.
54    pin_memory : bool, default False
55        Passed to ``DataLoader``.
56
57    Returns
58    -------
59    torch.utils.data.DataLoader
60    """
61    try:
62        import torch  # noqa: F401
63        from torch.utils.data import DataLoader, IterableDataset
64    except ImportError as e:  # pragma: no cover
65        raise ImportError(
66            "to_torch_dataloader requires torch. "
67            "Install with: pip install 'rpx-benchmark[depth-hf]' "
68            "(or any extra that pulls torch)."
69        ) from e
70
71    class _RPXIterableDataset(IterableDataset):
72        def __init__(self, source: Iterable[Any]) -> None:
73            self._source = source
74
75        def __iter__(self):  # type: ignore[override]
76            for batch in self._source:
77                yield batch
78
79    return DataLoader(
80        _RPXIterableDataset(rpx_iterable),
81        batch_size=None,  # the source already yields batches
82        num_workers=num_workers,
83        pin_memory=pin_memory,
84        collate_fn=rpx_collate,
85    )