maxent_grpo.config.dataset

Dataset and script argument dataclasses shared across training entrypoints.

These helpers extend TRL’s ScriptArguments to support dataset mixtures, including per-dataset column selection, sampling weights, and optional train/test splits. TRL is treated as an optional dependency so the config objects remain importable in lightweight environments such as doc builds and unit tests.

License Copyright 2025 Liv d’Aliberti

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

Classes

DatasetConfig(id[, config, split, columns, ...])

Configuration for a dataset inside a mixture.

DatasetMixtureConfig(datasets[, seed, ...])

Configuration for a mixture of datasets.

ScriptArguments([dataset_name, ...])

Extended TRL ScriptArguments with dataset mixture support.

_DatasetMixtureMeta

Metaclass that treats compatible mixtures as instances across reloads.

_ScriptArgumentsBase()

Fallback base for TRL ScriptArguments when type info is unavailable.
class maxent_grpo.config.dataset.DatasetConfig(id, config=None, split='train', columns=None, weight=None)[source]

Bases: object

Configuration for a dataset inside a mixture.

This describes one dataset entry pulled from the Hub and optionally filtered/weighted before mixing.

Variables:

  • id – Dataset repository ID on the Hub (e.g., “org/name”).

  • config – Optional dataset configuration name.

  • split – Split to load, defaults to “train”.

  • columns – Optional list of column names to keep; if provided all datasets in the mixture must share the same set of columns.

  • weight – Optional sampling weight in (0, 1]; when specified, a subsample of that proportion is taken from the split.

Parameters:
id: str
config: str | None = None
split: str = 'train'
columns: list[str] | None = None
weight: float | None = None
class maxent_grpo.config.dataset.DatasetMixtureConfig(datasets, seed=0, test_split_size=None)[source]

Bases: object

Configuration for a mixture of datasets.

Variables:

  • datasets – Ordered dataset entries combined into a single iterable.

  • seed – Seed used for deterministic shuffling and sampling.

  • test_split_size – Optional fraction moved to a held-out test split.

Raises:

ValueError – If test_split_size is provided outside (0, 1).

Parameters:
datasets: List[DatasetConfig]
seed: int = 0
test_split_size: float | None = None
class maxent_grpo.config.dataset.ScriptArguments(dataset_name=None, dataset_mixture=None, dataset_config=None)[source]

Bases: ScriptArguments

Extended TRL ScriptArguments with dataset mixture support.

Accepts either a single dataset via dataset_name or a declarative mixture provided as a mapping. When a dictionary is supplied the payload is converted into DatasetMixtureConfig entries and validated for consistent column naming across datasets.

Variables:

  • dataset_name – Dataset ID on the Hub when using a single source.

  • dataset_mixture – Mixture configuration with constituent datasets.

  • dataset_config – Optional config name associated with dataset_name.

Raises:

ValueError – If neither a dataset nor mixture is provided, or if the mixture payload is malformed.

Parameters:
dataset_name: str | None = None
dataset_mixture: DatasetMixtureConfig | None = None
dataset_config: str | None = None