diff --git a/tests/test_cli_utils.py b/tests/test_cli_utils.py index c3eeff391fa..5e98e353d6b 100644 --- a/tests/test_cli_utils.py +++ b/tests/test_cli_utils.py @@ -385,6 +385,24 @@ def test_dataset_fraction_partial_raises_error(self): with pytest.raises(ValueError, match="must be set for either all datasets"): get_dataset(mixture_config) + @pytest.mark.parametrize( + ("fractions", "error_message"), + [ + ([0.0, 0.0], "must be greater than zero"), + ([-0.5, 1.5], "must be non-negative"), + ([1.5, -0.5], "must be non-negative"), + ], + ) + def test_dataset_fraction_invalid_raises_error_before_loading(self, fractions, error_message): + mixture_config = DatasetMixtureConfig( + datasets=[DatasetConfig(path="unused", fraction=fraction) for fraction in fractions] + ) + + with patch("datasets.load_dataset") as load_dataset_mock, pytest.raises(ValueError, match=error_message): + get_dataset(mixture_config) + + load_dataset_mock.assert_not_called() + def test_dataset_fraction_streaming_raises_error(self): mixture_config = DatasetMixtureConfig( datasets=[DatasetConfig(path="trl-internal-testing/zen", name="standard_language_modeling", fraction=0.5)], diff --git a/trl/scripts/utils.py b/trl/scripts/utils.py index 89919961e20..fd65a3da329 100644 --- a/trl/scripts/utils.py +++ b/trl/scripts/utils.py @@ -65,7 +65,7 @@ class DatasetConfig: datasets, and the mixture size is capped so that no dataset is oversampled: the first `round(fraction * N)` rows of each dataset are kept, where `N` is the largest mixture size that avoids oversampling. Must be set for either all datasets in the mixture or none of them. Not supported for streaming datasets. When unset, - the full datasets are concatenated. + the full datasets are concatenated. Values must be non-negative and their sum must be greater than zero. """ path: str @@ -442,6 +442,20 @@ def get_dataset(mixture_config: DatasetMixtureConfig) -> "DatasetDict": import datasets logger.info(f"Creating dataset mixture with {len(mixture_config.datasets)} datasets") + + # Validate fractions before loading any datasets so invalid configurations fail with an actionable error. + fractions = [dataset_config.fraction for dataset_config in mixture_config.datasets] + if any(fraction is not None for fraction in fractions): + if any(fraction is None for fraction in fractions): + raise ValueError("`fraction` must be set for either all datasets in the mixture or none of them.") + if any(fraction < 0 for fraction in fractions): + raise ValueError("Dataset `fraction` values must be non-negative.") + total_fraction = sum(fractions) + if total_fraction <= 0: + raise ValueError("The sum of dataset `fraction` values must be greater than zero.") + if mixture_config.streaming: + raise ValueError("Using a dataset `fraction` is not supported with streaming datasets.") + datasets_list = [] for dataset_config in mixture_config.datasets: logger.info(f"Loading dataset for mixture: {dataset_config.path} (config name: {dataset_config.name})") @@ -460,13 +474,8 @@ def get_dataset(mixture_config: DatasetMixtureConfig) -> "DatasetDict": # If `fraction` is set, treat the values as target shares of the final mixture. They are normalized to sum to one, # and the mixture size is capped so that no dataset is oversampled: we keep the first `round(weight * total)` rows of # each dataset, where `total` is the largest mixture size such that no dataset contributes more rows than it has. - fractions = [dataset_config.fraction for dataset_config in mixture_config.datasets] if any(fraction is not None for fraction in fractions): - if any(fraction is None for fraction in fractions): - raise ValueError("`fraction` must be set for either all datasets in the mixture or none of them.") - if mixture_config.streaming: - raise ValueError("Using a dataset `fraction` is not supported with streaming datasets.") - weights = [fraction / sum(fractions) for fraction in fractions] + weights = [fraction / total_fraction for fraction in fractions] total = min( len(dataset) / weight for dataset, weight in zip(datasets_list, weights, strict=False) if weight > 0 )