omniloader 1.0.0__py3-none-any.whl

omniloader/config.py

@@ -0,0 +1,457 @@
+"""Configuration and reproducibility helpers.
+
+:class:`OmniConfig` bundles the knobs a training run needs — seed, batch size,
+worker count, the mixing strategy and the normalization/augmentation pipeline —
+and loads from a dict, a JSON file or a YAML file so experiments stay
+declarative. :func:`seed_everything` seeds Python, NumPy and PyTorch in one call.
+"""
+
+from __future__ import annotations
+
+import json
+import random
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+import yaml
+
+from omniloader.sampling.strategies import (
+    AnnealedTemperatureStrategy,
+    FixedWeightStrategy,
+    MixingStrategy,
+    ProportionalStrategy,
+    RoundRobinStrategy,
+    SubsampleConfig,
+    TemperatureStrategy,
+)
+from omniloader.sampling.subsamplers import ExhaustionPolicy
+from omniloader.transforms import build_transform
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from torch.utils.data import DataLoader
+
+    from omniloader.loader import OmniLoader, SizedDataset
+    from omniloader.schema.spec import DatasetSchema, UnifiedSchema
+    from omniloader.transforms import Compose
+
+#: Registry mapping strategy names to their classes.
+STRATEGIES: dict[str, type[MixingStrategy]] = {
+    "proportional": ProportionalStrategy,
+    "temperature": TemperatureStrategy,
+    "annealed_temperature": AnnealedTemperatureStrategy,
+    "fixed": FixedWeightStrategy,
+    "round_robin": RoundRobinStrategy,
+}
+
+#: Recognised ``collate`` names for :meth:`OmniConfig.build_collate`.
+COLLATES: tuple[str, ...] = ("unified", "dynamic", "mixup")
+
+
+def _subsample_from_dict(entry: dict[str, Any]) -> SubsampleConfig:
+    """Build a :class:`SubsampleConfig` from a config dict.
+
+    Args:
+        entry: Mapping with any of ``replacement``, ``policy`` (``"fresh"`` or
+            ``"exhaust"``), ``effective_size`` and ``sample_weights``.
+
+    Returns:
+        The populated :class:`SubsampleConfig`.
+
+    """
+    kwargs: dict[str, Any] = dict(entry)
+    if "policy" in kwargs and not isinstance(kwargs["policy"], ExhaustionPolicy):
+        kwargs["policy"] = ExhaustionPolicy(kwargs["policy"])
+    return SubsampleConfig(**kwargs)
+
+
+def seed_everything(seed: int) -> int:
+    """Seed Python, NumPy and PyTorch RNGs for reproducibility.
+
+    Args:
+        seed: The seed value.
+
+    Returns:
+        The seed, for convenience.
+
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    return seed
+
+
+@dataclass
+class OmniConfig:
+    """Declarative configuration for an OmniLoader training run.
+
+    Args:
+        seed: Global RNG seed.
+        batch_size: Batch size for the ``DataLoader``.
+        num_workers: Number of worker processes for the ``DataLoader``.
+        pad_features: Whether the unifier pads/crops sequence values.
+        strategy: Name of the mixing strategy (a key of :data:`STRATEGIES`).
+        strategy_kwargs: Extra keyword arguments forwarded to the strategy
+            (e.g. ``temperature``, ``weights``, ``samples_per_dataset``).
+        transforms: Ordered list of transform config dicts (each with a ``name``
+            key) applied during training. Stats-based normalizers may reference a
+            saved stats file via ``stats_path`` (or ``union_stats_path``) instead of
+            inlining ``stats``. See :func:`~omniloader.transforms.build_transform`.
+        eval_transforms: Transforms applied during evaluation. When ``None``,
+            only the non-augmenting (train-agnostic) transforms in
+            :attr:`transforms`, such as ``normalize``, take effect.
+        datasets: Declarative dataset entries (adapter + args + schema) built by
+            :meth:`build_datasets`. See :mod:`omniloader.data.factory`.
+        subsample: Optional per-dataset subsampling entries (one per dataset, or
+            ``None`` to skip a dataset), each a dict for :class:`SubsampleConfig`.
+            Built by :meth:`build_subsample` and fed to :meth:`build_strategy`.
+        collate: Collate function name — one of :data:`COLLATES` (``"unified"``,
+            ``"dynamic"`` or ``"mixup"``). See :meth:`build_collate`. ``"dynamic"``
+            only reduces padding with :attr:`pad_features` set to ``False`` (under a
+            fixed ``time_dim`` it behaves like ``"unified"``).
+        collate_kwargs: Extra keyword arguments for the collator. For ``"dynamic"``
+            an optional ``keys`` list; for ``"mixup"`` a ``base`` collate name plus
+            ``alpha``/``mode``/``p``.
+        bucketing: Optional length-bucketing entry with a ``key`` (sequence name to
+            bucket by) and optional ``batch_size``/``bucket_multiplier``/
+            ``drop_last``/``shuffle``. When set, :meth:`build_dataloader` batches
+            with :class:`~omniloader.sampling.bucketing.LengthBucketBatchSampler`.
+            Only meaningful with ``pad_features=False`` (a no-op under a fixed
+            ``time_dim``).
+        num_replicas: Distributed replica count for the sampler. ``None`` infers
+            it from ``torch.distributed``.
+        rank: This process's distributed rank. ``None`` infers it.
+        drop_last: Drop the uneven tail when sharding across replicas.
+        pin_memory: Copy tensors into pinned memory before returning them (faster
+            host-to-GPU transfer). Forwarded to the ``DataLoader``.
+        persistent_workers: Keep worker processes alive between epochs. Only takes
+            effect when :attr:`num_workers` > 0.
+        prefetch_factor: Batches prefetched per worker. Only applied when
+            :attr:`num_workers` > 0; ``None`` uses the PyTorch default.
+
+    """
+
+    seed: int = 0
+    batch_size: int = 32
+    num_workers: int = 0
+    pad_features: bool = True
+    strategy: str = "proportional"
+    strategy_kwargs: dict[str, Any] = field(default_factory=dict)
+    transforms: list[dict[str, Any]] = field(default_factory=list)
+    eval_transforms: list[dict[str, Any]] | None = None
+    datasets: list[dict[str, Any]] = field(default_factory=list)
+    subsample: list[dict[str, Any] | None] | None = None
+    collate: str = "unified"
+    collate_kwargs: dict[str, Any] = field(default_factory=dict)
+    bucketing: dict[str, Any] | None = None
+    num_replicas: int | None = None
+    rank: int | None = None
+    drop_last: bool = False
+    pin_memory: bool = False
+    persistent_workers: bool = False
+    prefetch_factor: int | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> OmniConfig:
+        """Build a config from a plain dict, ignoring unknown keys.
+
+        Args:
+            data: Mapping of field names to values.
+
+        Returns:
+            The populated config.
+
+        """
+        fields = set(cls.__dataclass_fields__)
+        return cls(**{k: v for k, v in data.items() if k in fields})
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> OmniConfig:
+        """Load a config from a YAML file."""
+        with Path(path).open("r", encoding="utf-8") as f:
+            return cls.from_dict(yaml.safe_load(f) or {})
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> OmniConfig:
+        """Load a config from a JSON file."""
+        with Path(path).open("r", encoding="utf-8") as f:
+            return cls.from_dict(json.load(f) or {})
+
+    @classmethod
+    def from_file(cls, path: str | Path) -> OmniConfig:
+        """Load a config from a ``.json``, ``.yaml`` or ``.yml`` file by extension.
+
+        Args:
+            path: Path to a config file.
+
+        Returns:
+            The populated config.
+
+        Raises:
+            ValueError: If the file extension is not recognised.
+
+        """
+        suffix = Path(path).suffix.lower()
+        if suffix == ".json":
+            return cls.from_json(path)
+        if suffix in {".yaml", ".yml"}:
+            return cls.from_yaml(path)
+        raise ValueError(f"Unsupported config extension {suffix!r}; use .json/.yaml/.yml")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return the config as a plain dict."""
+        return {
+            "seed": self.seed,
+            "batch_size": self.batch_size,
+            "num_workers": self.num_workers,
+            "pad_features": self.pad_features,
+            "strategy": self.strategy,
+            "strategy_kwargs": dict(self.strategy_kwargs),
+            "transforms": [dict(t) for t in self.transforms],
+            "eval_transforms": (
+                None if self.eval_transforms is None else [dict(t) for t in self.eval_transforms]
+            ),
+            "datasets": [dict(d) for d in self.datasets],
+            "subsample": (
+                None
+                if self.subsample is None
+                else [None if s is None else dict(s) for s in self.subsample]
+            ),
+            "collate": self.collate,
+            "collate_kwargs": dict(self.collate_kwargs),
+            "bucketing": None if self.bucketing is None else dict(self.bucketing),
+            "num_replicas": self.num_replicas,
+            "rank": self.rank,
+            "drop_last": self.drop_last,
+            "pin_memory": self.pin_memory,
+            "persistent_workers": self.persistent_workers,
+            "prefetch_factor": self.prefetch_factor,
+        }
+
+    def build_strategy(
+        self,
+        sizes: Sequence[int],
+        subsample: Sequence[SubsampleConfig | None] | None = None,
+    ) -> MixingStrategy:
+        """Instantiate the configured mixing strategy for the given sizes.
+
+        Args:
+            sizes: Per-dataset sample counts.
+            subsample: Optional per-dataset subsample configs.
+
+        Returns:
+            The configured :class:`~omniloader.sampling.strategies.MixingStrategy`.
+
+        Raises:
+            ValueError: If :attr:`strategy` is not a known strategy name.
+
+        """
+        if self.strategy not in STRATEGIES:
+            raise ValueError(
+                f"Unknown strategy {self.strategy!r}; expected one of {sorted(STRATEGIES)}"
+            )
+        strategy_cls = STRATEGIES[self.strategy]
+        return strategy_cls(sizes, seed=self.seed, subsample=subsample, **self.strategy_kwargs)
+
+    def build_transform(self, schema: UnifiedSchema, training: bool = True) -> Compose | None:
+        """Build the transform pipeline for a split.
+
+        Args:
+            schema: The unified schema, injected into schema-aware transforms.
+            training: When ``True``, uses :attr:`transforms`; otherwise uses
+                :attr:`eval_transforms` if set, else :attr:`transforms` (whose
+                augmentations self-skip in eval mode).
+
+        Returns:
+            A :class:`~omniloader.transforms.Compose`, or ``None`` if empty.
+
+        """
+        configs = self.transforms
+        if not training and self.eval_transforms is not None:
+            configs = self.eval_transforms
+        return build_transform(configs, schema)
+
+    def build_datasets(self) -> tuple[list[SizedDataset], list[DatasetSchema]]:
+        """Construct the datasets/schemas declared in :attr:`datasets`.
+
+        Returns:
+            A ``(datasets, schemas)`` tuple. See
+            :func:`omniloader.data.factory.build_datasets`.
+
+        """
+        from omniloader.data.factory import build_datasets
+
+        return build_datasets(self.datasets)
+
+    def build_subsample(self) -> list[SubsampleConfig | None] | None:
+        """Build the per-dataset :class:`SubsampleConfig` list from :attr:`subsample`.
+
+        Returns:
+            One :class:`SubsampleConfig` (or ``None``) per dataset, or ``None`` when
+            :attr:`subsample` is unset.
+
+        """
+        if self.subsample is None:
+            return None
+        return [None if entry is None else _subsample_from_dict(entry) for entry in self.subsample]
+
+    def build_collate(
+        self, schema: UnifiedSchema, training: bool = True
+    ) -> Callable[[Sequence[dict[str, Any]]], dict[str, Any]]:
+        """Build the collate function named by :attr:`collate`.
+
+        Args:
+            schema: The unified schema, injected into schema-aware collators.
+            training: When ``False``, the batch-level ``"mixup"`` augmentation is
+                disabled and its ``base`` collator is used instead.
+
+        Returns:
+            A callable turning a list of samples into a batched dict.
+
+        Raises:
+            ValueError: If :attr:`collate` (or a ``mixup`` ``base``) is unknown.
+
+        """
+        from omniloader.collate import DynamicCollator, unified_collate
+
+        if self.collate not in COLLATES:
+            raise ValueError(f"Unknown collate {self.collate!r}; expected one of {list(COLLATES)}")
+
+        def _base(name: str, kwargs: dict[str, Any]) -> Any:
+            if name == "unified":
+                return unified_collate
+            if name == "dynamic":
+                return DynamicCollator(schema, keys=kwargs.get("keys"))
+            raise ValueError(f"Unknown collate {name!r}; expected 'unified' or 'dynamic'")
+
+        if self.collate != "mixup":
+            return _base(self.collate, self.collate_kwargs)
+
+        from omniloader.transforms.mix import MixupCollator
+
+        kwargs = dict(self.collate_kwargs)
+        base = _base(kwargs.pop("base", "unified"), kwargs)
+        kwargs.pop("keys", None)  # consumed by the base collator only
+        if not training:
+            return base  # mixing is a train-time augmentation
+        kwargs.setdefault("seed", self.seed)
+        return MixupCollator(base, schema, **kwargs)
+
+    def build_loader(
+        self,
+        datasets: Sequence[SizedDataset] | None = None,
+        schemas: Sequence[DatasetSchema] | None = None,
+        training: bool = True,
+    ) -> OmniLoader:
+        """Build a fully configured :class:`~omniloader.loader.OmniLoader`.
+
+        Args:
+            datasets: Source datasets. When ``None``, they are built from
+                :attr:`datasets` via :meth:`build_datasets`.
+            schemas: Per-dataset schemas. Required when ``datasets`` is given;
+                ignored (rebuilt) when ``datasets`` is ``None``.
+            training: Whether augmentations are active (``False`` for val/test).
+
+        Returns:
+            The loader with its transform pipeline wired in.
+
+        Raises:
+            ValueError: If ``datasets`` is given without ``schemas``.
+
+        """
+        from omniloader.loader import OmniLoader
+
+        if datasets is None:
+            datasets, schemas = self.build_datasets()
+        elif schemas is None:
+            raise ValueError("schemas must be given when datasets is provided")
+        loader = OmniLoader(
+            datasets,
+            schemas,
+            pad_features=self.pad_features,
+            training=training,
+            seed=self.seed,
+        )
+        loader.transform = self.build_transform(loader.schema, training=training)
+        return loader
+
+    def build_dataloader(
+        self,
+        datasets: Sequence[SizedDataset] | None = None,
+        schemas: Sequence[DatasetSchema] | None = None,
+        training: bool = True,
+        epoch: int = 0,
+    ) -> DataLoader:
+        """Assemble a ready-to-iterate ``DataLoader`` entirely from this config.
+
+        Wires together the loader, mixing strategy, distributed sampler, collate
+        function and (optionally) length bucketing. Training loaders draw indices
+        through the configured :class:`~omniloader.sampling.strategies.MixingStrategy`;
+        evaluation loaders (``training=False``) iterate sequentially for a stable,
+        full-coverage pass.
+
+        Args:
+            datasets: Source datasets, or ``None`` to build them from
+                :attr:`datasets`.
+            schemas: Per-dataset schemas (required with ``datasets``).
+            training: Whether this is the training split.
+            epoch: Initial epoch for the sampler and per-sample augmentation.
+
+        Returns:
+            A configured :class:`torch.utils.data.DataLoader`.
+
+        """
+        from torch.utils.data import DataLoader, SequentialSampler
+
+        from omniloader.sampling.sampler import OmniSampler
+        from omniloader.utils.seeding import seed_worker
+
+        loader = self.build_loader(datasets, schemas, training=training)
+        loader.set_epoch(epoch)
+        collate = self.build_collate(loader.schema, training=training)
+
+        sampler: OmniSampler | None = None
+        if training:
+            strategy = self.build_strategy(loader.dataset_sizes, subsample=self.build_subsample())
+            sampler = OmniSampler(
+                strategy,
+                epoch=epoch,
+                num_replicas=self.num_replicas,
+                rank=self.rank,
+                drop_last=self.drop_last,
+            )
+
+        # persistent_workers/prefetch_factor are only valid with worker processes.
+        kwargs: dict[str, Any] = {
+            "num_workers": self.num_workers,
+            "collate_fn": collate,
+            "worker_init_fn": seed_worker,
+            "pin_memory": self.pin_memory,
+        }
+        if self.num_workers > 0:
+            kwargs["persistent_workers"] = self.persistent_workers
+            if self.prefetch_factor is not None:
+                kwargs["prefetch_factor"] = self.prefetch_factor
+
+        if self.bucketing is not None:
+            from omniloader.sampling.bucketing import LengthBucketBatchSampler
+
+            opts = dict(self.bucketing)
+            key = opts.pop("key")
+            batch_size = opts.pop("batch_size", self.batch_size)
+            base_sampler = sampler if sampler is not None else SequentialSampler(loader)
+            opts.setdefault("shuffle", training)
+            batch_sampler = LengthBucketBatchSampler(
+                base_sampler,
+                loader.sequence_lengths(key),
+                batch_size,
+                seed=self.seed,
+                **opts,
+            )
+            return DataLoader(loader, batch_sampler=batch_sampler, **kwargs)
+
+        return DataLoader(loader, batch_size=self.batch_size, sampler=sampler, **kwargs)

omniloader/data/datasets.py

@@ -0,0 +1,196 @@
+"""Dataset adapters that yield raw sample dicts for OmniLoader.
+
+The concrete datasets here (an HDF5-backed reader and an in-memory tensor
+dataset) implement the loose contract OmniLoader expects: ``__len__`` returns the
+number of samples and ``__getitem__`` returns a ``dict`` of named values. The
+:class:`~omniloader.schema.spec.DatasetSchema` describing what a dataset provides is
+supplied alongside it to :class:`~omniloader.loader.OmniLoader`.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+
+def _decode_hdf5_value(dset: h5py.Dataset) -> Any:
+    """Read a single HDF5 dataset into a tensor, string, or list of strings.
+
+    Args:
+        dset: An open HDF5 dataset object.
+
+    Returns:
+        A :class:`torch.Tensor` for numeric data, a ``str`` for scalar strings,
+        or a ``list[str]`` for string arrays.
+
+    """
+    if dset.dtype.kind in {"S", "O"}:
+        value = dset.asstr()[()]
+        if isinstance(value, (np.ndarray, list)):
+            return [str(v) for v in value]
+        return str(value)
+
+    value = dset[()]
+    if isinstance(value, np.ndarray):
+        return torch.from_numpy(value)
+    if isinstance(value, (bool, np.bool_)):
+        return torch.tensor(bool(value))
+    if isinstance(value, (int, np.integer)):
+        return torch.tensor(int(value))
+    if isinstance(value, (float, np.floating)):
+        return torch.tensor(float(value))
+    return torch.as_tensor(value)
+
+
+class HDF5Dataset(Dataset):
+    """Read samples from a group of an HDF5 file, one sub-group per sample.
+
+    The file is expected to hold a group named ``subset`` (e.g. ``"train"``)
+    whose children are per-sample groups, each containing the named value
+    datasets (features, targets, masks, metadata). The file handle is opened
+    lazily per worker process so the dataset is safe to use with
+    ``num_workers > 0``.
+
+    Args:
+        h5_path: Path to the ``.h5`` file.
+        subset: Name of the top-level group to read (e.g. ``"train"``).
+        keys: Optional subset of value names to load per sample. When ``None``,
+            every dataset in the sample group is loaded.
+        cache_size: If ``> 0``, keep an in-process LRU cache of up to this many
+            decoded samples to speed up repeated reads.
+        preload: If ``True``, decode every sample into memory up front (fastest
+            for small datasets; uses the most memory).
+
+    Raises:
+        KeyError: If ``subset`` is not present in the file.
+
+    """
+
+    def __init__(
+        self,
+        h5_path: str | Path,
+        subset: str,
+        keys: Sequence[str] | None = None,
+        cache_size: int = 0,
+        preload: bool = False,
+    ) -> None:
+        self.h5_path = Path(h5_path)
+        self.subset = subset
+        self.keys = list(keys) if keys is not None else None
+        self.cache_size = cache_size
+        self._file: h5py.File | None = None
+        self._cache: OrderedDict[int, dict[str, Any]] = OrderedDict()
+        with h5py.File(self.h5_path, "r") as f:
+            if subset not in f:
+                raise KeyError(f"Subset {subset!r} not found in {self.h5_path}")
+            self.sample_ids = sorted(f[subset].keys())
+        self._preloaded: list[dict[str, Any]] | None = (
+            [self._read(i) for i in range(len(self.sample_ids))] if preload else None
+        )
+
+    def _handle(self) -> h5py.File:
+        """Return a per-process open file handle, opening it on first use."""
+        if self._file is None:
+            self._file = h5py.File(self.h5_path, "r")
+        return self._file
+
+    def _read(self, index: int) -> dict[str, Any]:
+        """Decode the sample at ``index`` directly from the HDF5 file."""
+        sample_id = self.sample_ids[index]
+        group = self._handle()[self.subset][sample_id]
+        names = self.keys if self.keys is not None else list(group.keys())
+        return {name: _decode_hdf5_value(group[name]) for name in names if name in group}
+
+    def __len__(self) -> int:
+        """Number of samples in the subset."""
+        return len(self.sample_ids)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        """Load the sample at ``index`` as a dict of named values.
+
+        Args:
+            index: Position within the subset.
+
+        Returns:
+            A dict mapping each stored key to a tensor or string.
+
+        """
+        if self._preloaded is not None:
+            return self._preloaded[index]
+        if self.cache_size > 0 and index in self._cache:
+            self._cache.move_to_end(index)  # mark as most-recently used
+            return self._cache[index]
+        sample = self._read(index)
+        if self.cache_size > 0:
+            self._cache[index] = sample
+            if len(self._cache) > self.cache_size:
+                self._cache.popitem(last=False)  # evict least-recently used
+        return sample
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Drop the open file handle and cache when pickling (for workers)."""
+        state = self.__dict__.copy()
+        state["_file"] = None
+        state["_cache"] = OrderedDict()
+        return state
+
+
+class DictTensorDataset(Dataset):
+    """In-memory dataset over a dict of stacked tensors, keyed by sample index.
+
+    Useful for tests and small experiments: each entry maps a name to a tensor
+    whose leading dimension is the sample axis. ``__getitem__`` returns the
+    per-sample slice of every tensor as a dict.
+
+    Args:
+        tensors: Mapping of value name to a tensor of shape ``(N, ...)`` where
+            ``N`` is the number of samples (equal across all tensors).
+        metadata: Optional mapping of metadata name to a per-sample sequence
+            (e.g. dataset name or sample key) of length ``N``.
+
+    Raises:
+        ValueError: If the tensors disagree on the number of samples.
+
+    """
+
+    def __init__(
+        self,
+        tensors: Mapping[str, torch.Tensor],
+        metadata: Mapping[str, Sequence[Any]] | None = None,
+    ) -> None:
+        if not tensors:
+            raise ValueError("DictTensorDataset requires at least one tensor")
+        lengths = {t.shape[0] for t in tensors.values()}
+        if len(lengths) != 1:
+            raise ValueError(f"All tensors must share the sample dimension, got {lengths}")
+        self.tensors = dict(tensors)
+        self.metadata = {k: list(v) for k, v in (metadata or {}).items()}
+        self._length = lengths.pop()
+
+    def __len__(self) -> int:
+        """Number of samples."""
+        return self._length
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        """Return the per-sample slice of every tensor plus metadata.
+
+        Args:
+            index: Sample position.
+
+        Returns:
+            A dict of the sample's values and metadata.
+
+        """
+        sample: dict[str, Any] = {name: t[index] for name, t in self.tensors.items()}
+        for name, values in self.metadata.items():
+            sample[name] = values[index]
+        return sample