def test_mixture_column_selection(self):
        """Test loading a dataset mixture with column selection."""
        dataset_configs = [
            DatasetConfig(
                id=self.dataset_name,
                config=self.dataset_config,
                split="train",
                columns=["prompt", "chosen"],
                weight=None,
            ),
        ]
        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.assertIn("prompt", dataset["train"].column_names)
        self.assertIn("chosen", dataset["train"].column_names)
Test loading a dataset mixture with column selection.
tests/utils/test_data.py