omniloader.loader

The OmniLoader meta-dataset that unifies several disjoint datasets.

OmniLoader concatenates any number of source datasets, each paired with a DatasetSchema, and presents them as one dataset whose every sample follows the shared UnifiedSchema. Keys a given source lacks are filled with placeholder tensors and all-False masks by the SampleUnifier, so a single batch may freely mix samples from different datasets. A MixingStrategy (obtained via OmniLoader.make_sampler()) controls how the datasets are balanced across an epoch.

Classes

SizedDataset	A map-style dataset that reports its length and returns dict samples.
OmniLoader	Unify multiple datasets into a single masked, concatenated dataset.

Module Contents

class omniloader.loader.SizedDataset

Bases: Protocol

A map-style dataset that reports its length and returns dict samples.

__len__() → int

Return the number of samples.

__getitem__(index: int) → dict[str, Any]

Return the sample at index.

class omniloader.loader.OmniLoader(datasets: collections.abc.Sequence[SizedDataset], schemas: collections.abc.Sequence[omniloader.schema.spec.DatasetSchema], schema: omniloader.schema.spec.UnifiedSchema | None = None, pad_features: bool = True, transform: omniloader.transforms.Transform | None = None, training: bool = True, seed: int = 0)

Bases: torch.utils.data.Dataset

Unify multiple datasets into a single masked, concatenated dataset.

Parameters:

• datasets – The source datasets. Each must implement __len__ and return a dict from __getitem__.

• schemas – One DatasetSchema per dataset, declaring what it provides. Merged into a UnifiedSchema when schema is not given.

• schema – Optional pre-built unified schema. When None, it is the union of schemas.

• pad_features – Whether the unifier pads/crops sequence values to their declared time_dim. See SampleUnifier.

• transform – Optional per-sample transform (normalization/augmentation) applied after unification. See omniloader.transforms.

• training – Whether augmentations in transform are active. Set False for validation/test loaders.

• seed – Base seed for per-sample augmentation randomness. Combined with the current epoch (see set_epoch()) and the sample index so that augmentation is reproducible and independent of worker count.

Raises:

ValueError – If datasets is empty or its length differs from schemas.

set_epoch(epoch: int) → None

Set the epoch used to seed per-sample augmentation.

Call once per epoch (the OmniDataModule does this automatically) so augmentation differs across epochs yet stays reproducible.

Parameters:

epoch – The upcoming epoch number.

__len__() → int

Total number of samples across all source datasets.

locate(index: int) → tuple[int, int]

Map a global index to (dataset_index, local_index).

Parameters:

index – Global index in [0, len(self)). Negative indices count from the end.

Returns:

The owning dataset’s position and the local index within it.

Raises:

IndexError – If index is out of range.

__getitem__(index: int) → dict[str, Any]

Return the unified sample at a global index.

Parameters:

index – Global index into the concatenated dataset.

Returns:

The sample mapped onto the unified schema with masks and metadata.

describe(max_samples: int = 64) → omniloader.introspection.Report

Summarise this loader’s datasets (coverage, valid fractions, classes).

Parameters:

max_samples – Samples read per dataset for the statistics.

Returns:

A Report.

validate(num_samples: int = 4, strict: bool = False) → list[str]

Check each dataset’s raw samples against its declared specs.

Parameters:

• num_samples – Samples probed per dataset.

• strict – Raise ValueError if any mismatch is found.

Returns:

A list of issue strings (empty when everything matches).

sequence_lengths(key: str) → list[int]

Return the native sequence length of key for every sample.

Reads each source sample’s raw value (bypassing unification/transforms) and reports the length of its sequence axis, or 0 when the sample lacks the key. Intended to feed LengthBucketBatchSampler.

Note

This iterates the entire dataset once — O(N) reads. For HDF5 back ends prefer computing it once and caching the result.

Parameters:

key – A sequence feature or target name.

Returns:

A list of per-sample lengths, indexed by global sample index.

make_sampler(strategy: omniloader.sampling.strategies.MixingStrategy | None = None, seed: int = 0, subsample: collections.abc.Sequence[omniloader.sampling.strategies.SubsampleConfig | None] | None = None, epoch: int = 0) → omniloader.sampling.sampler.OmniSampler

Build an OmniSampler for driving a DataLoader.

Parameters:

• strategy – A mixing strategy. When None, a ProportionalStrategy over the dataset sizes is used.

• seed – Base seed for the default strategy (ignored if strategy is given).

• subsample – Per-dataset subsample configs for the default strategy.

• epoch – Initial epoch for the sampler.

Returns:

A sampler yielding global indices per epoch.