def test_weighted_mixture(self):
        """Test loading a dataset mixture with weights."""
        dataset_configs = [
            DatasetConfig(id=self.dataset_name, config=self.dataset_config, split="train", columns=None, weight=0.25),
            DatasetConfig(id=self.dataset_name, config=self.dataset_config, split="test", columns=None, weight=0.5),
        ]
        dataset_mixture = DatasetMixtureConfig(
            datasets=dataset_configs,
        )
        args = ScriptArguments(dataset_mixture=asdict(dataset_mixture))
        dataset = get_dataset(args)
        self.assertIsInstance(dataset, DatasetDict)
        self.assertIn("train", dataset)
        self.assertEqual(
            len(dataset["train"]), len(self.ref_dataset["train"]) // 4 + len(self.ref_dataset["test"]) // 2
        )
Test loading a dataset mixture with weights.
tests/utils/test_data.py