omniloader.sampling.strategies

Dataset-mixing strategies that decide the per-epoch sampling scheme.

Each strategy answers one question: given several datasets of possibly very different sizes, how many samples from each should the network see per epoch, and in what order? Strategies return a flat list of global indices (into the concatenated OmniLoader dataset) via MixingStrategy.epoch_indices(), so a plain OmniSampler can drive an ordinary DataLoader.

Available strategies:

• ProportionalStrategy — sample in proportion to true dataset sizes.

• TemperatureStrategy — re-weight by size ** (1 / temperature) (temperature=2 recovers square-root scaling).

• FixedWeightStrategy — explicit per-dataset weights, or uniform.

• RoundRobinStrategy — a fixed effective size per dataset, interleaved.

Classes

SubsampleConfig	Per-dataset sampling options.
MixingStrategy	Base class turning per-dataset sizes into per-epoch global index lists.
ProportionalStrategy	Sample each dataset in proportion to its true size (full concatenation).
TemperatureStrategy	Re-weight datasets by size ** (1 / temperature).
FixedWeightStrategy	Sample datasets according to explicit weights (or uniformly).
RoundRobinStrategy	Draw a fixed number of samples per dataset and interleave them.
AnnealedTemperatureStrategy	Temperature sampling whose temperature anneals across epochs.

Module Contents

class omniloader.sampling.strategies.SubsampleConfig

Per-dataset sampling options.

Parameters:

• replacement – Whether indices may repeat within an epoch.

• policy – Exhaustion policy for without-replacement draws.

• effective_size – Override for how many samples this dataset contributes per epoch. When None, the strategy decides.

• sample_weights – Optional per-sample weights (length = dataset size) for skewed within-dataset sampling, e.g. class balancing.

class omniloader.sampling.strategies.MixingStrategy(sizes: collections.abc.Sequence[int], seed: int = 0, subsample: collections.abc.Sequence[SubsampleConfig | None] | None = None)

Bases: abc.ABC

Base class turning per-dataset sizes into per-epoch global index lists.

Parameters:

• sizes – Number of samples in each dataset, in OmniLoader order.

• seed – Base seed for reproducible draws and shuffles.

• subsample – Optional per-dataset SubsampleConfig. Missing entries fall back to a default config.

Raises:

ValueError – If sizes is empty or contains a non-positive entry.

abstractmethod target_counts(epoch: int) → list[int]

Number of samples each dataset contributes in epoch.

Parameters:

epoch – The epoch number (used only by epoch-dependent strategies such as AnnealedTemperatureStrategy).

Returns:

A list with one non-negative count per dataset.

epoch_indices(epoch: int) → list[int]

Build the shuffled list of global indices to visit this epoch.

Parameters:

epoch – The epoch number, used to seed draws and the final shuffle.

Returns:

A flat list of global indices into the concatenated dataset.

__len__() → int

Total number of samples produced per epoch (evaluated at epoch 0).

class omniloader.sampling.strategies.ProportionalStrategy(sizes: collections.abc.Sequence[int], seed: int = 0, subsample: collections.abc.Sequence[SubsampleConfig | None] | None = None)

Bases: MixingStrategy

Sample each dataset in proportion to its true size (full concatenation).

target_counts(epoch: int) → list[int]

Return each dataset’s full size.

class omniloader.sampling.strategies.TemperatureStrategy(sizes: collections.abc.Sequence[int], temperature: float = 2.0, seed: int = 0, subsample: collections.abc.Sequence[SubsampleConfig | None] | None = None)

Bases: MixingStrategy

Re-weight datasets by size ** (1 / temperature).

temperature=1 reduces to proportional sampling; temperature=2 is square-root scaling; larger values push the mixture toward uniform. The total epoch size is preserved (equal to the sum of dataset sizes) unless per-dataset effective_size overrides are set.

Parameters:

• sizes – Per-dataset sizes.

• temperature – Sampling temperature >= 1.

• seed – Base seed.

• subsample – Optional per-dataset configs.

Raises:

ValueError – If temperature is less than 1.

target_counts(epoch: int) → list[int]

Return per-dataset counts weighted by size ** (1 / temperature).

class omniloader.sampling.strategies.FixedWeightStrategy(sizes: collections.abc.Sequence[int], weights: collections.abc.Sequence[float] | None = None, epoch_size: int | None = None, seed: int = 0, subsample: collections.abc.Sequence[SubsampleConfig | None] | None = None)

Bases: MixingStrategy

Sample datasets according to explicit weights (or uniformly).

Parameters:

• sizes – Per-dataset sizes.

• weights – Relative weight per dataset. When None, all datasets are weighted equally (uniform mixing). Need not sum to one.

• epoch_size – Total samples per epoch. Defaults to the sum of sizes.

• seed – Base seed.

• subsample – Optional per-dataset configs.

Raises:

ValueError – If weights has the wrong length or is not positive.

target_counts(epoch: int) → list[int]

Return per-dataset counts apportioned by the fixed weights.

class omniloader.sampling.strategies.RoundRobinStrategy(sizes: collections.abc.Sequence[int], samples_per_dataset: int | None = None, seed: int = 0, subsample: collections.abc.Sequence[SubsampleConfig | None] | None = None)

Bases: MixingStrategy

Draw a fixed number of samples per dataset and interleave them.

Unlike the weighted strategies, this one takes an explicit effective size per dataset (default: the smallest dataset’s size, so every dataset contributes equally) and interleaves the datasets round-robin rather than globally shuffling, which keeps batches diverse across datasets.

Parameters:

• sizes – Per-dataset sizes.

• samples_per_dataset – Number of samples to draw from each dataset each epoch. When None, uses the size of the smallest dataset.

• seed – Base seed.

• subsample – Optional per-dataset configs (effective_size overrides samples_per_dataset for that dataset).

target_counts(epoch: int) → list[int]

Return an equal count per dataset (smallest size unless overridden).

epoch_indices(epoch: int) → list[int]

Interleave datasets round-robin instead of globally shuffling.

Parameters:

epoch – The epoch number.

Returns:

Global indices ordered so consecutive samples come from different datasets where possible.

class omniloader.sampling.strategies.AnnealedTemperatureStrategy(sizes: collections.abc.Sequence[int], start_temperature: float = 5.0, end_temperature: float = 1.0, num_epochs: int = 10, seed: int = 0, subsample: collections.abc.Sequence[SubsampleConfig | None] | None = None)

Bases: MixingStrategy

Temperature sampling whose temperature anneals across epochs.

Multilingual/multi-task training often starts near-uniform (high temperature, up-weighting small datasets) and anneals toward proportional (temperature 1) so later epochs match the true data distribution. The temperature is linearly interpolated from start_temperature at epoch 0 to end_temperature at epoch num_epochs - 1 (and held there afterwards).

Parameters:

• sizes – Per-dataset sizes.

• start_temperature – Temperature at epoch 0 (>= 1).

• end_temperature – Temperature at the final epoch (>= 1).

• num_epochs – Number of epochs over which to anneal.

• seed – Base seed.

• subsample – Optional per-dataset configs.

Raises:

ValueError – If a temperature is below 1 or num_epochs is not positive.

temperature_at(epoch: int) → float

Return the interpolated temperature for epoch.

target_counts(epoch: int) → list[int]

Return counts using the epoch’s annealed temperature.