def test_interleave_datasets(dataset: IterableDataset, probas, seed, expected_length): d1 = dataset d2 = dataset.map(lambda x: {"id+1": x["id"] + 1, **x}) d3 = dataset.with_format("python") datasets = [d1, d2, d3] merged_dataset = interleave_datasets(datasets, probabilities=probas, seed=seed) # Check the examples iterable assert isinstance( merged_dataset._ex_iterable, (CyclingMultiSourcesExamplesIterable, RandomlyCyclingMultiSourcesExamplesIterable) ) # Check that it is deterministic if seed is not None: merged_dataset2 = interleave_datasets([d1, d2, d3], probabilities=probas, seed=seed) assert list(merged_dataset) == list(merged_dataset2) # Check first example if seed is not None: rng = np.random.default_rng(seed) i = next(iter(RandomlyCyclingMultiSourcesExamplesIterable._iter_random_indices(rng, len(datasets), p=probas))) assert next(iter(merged_dataset)) == next(iter(datasets[i])) else: assert any(next(iter(merged_dataset)) == next(iter(dataset)) for dataset in datasets) # Compute length it case it's random if expected_length is None: expected_length = 0 counts = [len(list(d)) for d in datasets] rng = np.random.default_rng(seed) for i in RandomlyCyclingMultiSourcesExamplesIterable._iter_random_indices(rng, len(datasets), p=probas): if counts[i] == 0: break counts[i] -= 1 expected_length += 1 # Check length assert len(list(merged_dataset)) == expected_length
def test_iterable_dataset_with_format(dataset: IterableDataset, format_type): formatted_dataset = dataset.with_format(format_type) assert formatted_dataset._format_type == format_type if format_type == "torch": import torch assert isinstance(formatted_dataset, torch.utils.data.IterableDataset)