Title: Feature proposal: Stacking, potentially heterogeneous, datasets
Author: TimCares (NONE)
State: open
Created: 2024-11-05T15:40:50Z

### Introduction

Hello there,
I noticed that there are two ways to combine multiple datasets: Either through `datasets.concatenate_datasets` or `datasets.interleave_datasets`. However, to my knowledge (please correct me if I am wrong) both approaches require the datasets that are combined to have the same features.

I think it would be a great idea to add support for combining multiple datasets that might not follow the same schema (i.e. have different features), for example an image and text dataset. That is why I propose a third function of the `datasets.combine` module called `stack_datasets`, which can be used to combine a list of datasets with (potentially) different features. This would look as follows:

```python
>>> from datasets import stack_datasets
>>> image_dataset = ...
>>> next(iter(image_dataset))
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=555x416 at 0x313E79CD0> }
>>> text_dataset = ...
>>> next(iter(text_dataset))
{'text': "This is a test."}
>>> stacked = stack_datasets(datasets={'i_ds': image_dataset, 't_ds': text_dataset}, stopping_strategy='all_exhausted')
>>> next(iter(stacked))
{
'i_ds': {'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=555x416 at 0x313E79CD0> }
't_ds': {'text': "This is a test."}
}
```
<br />

### Motivation

I motivate this by:

**A**: The fact that Pytorch offers a similar functionality under `torch.utils.data.StackDataset` ([link](https://pytorch.org/docs/stable/data.html#torch.utils.data.StackDataset)).

**B**: In settings where one would like to e.g. train a Vision-Language model using an image-text dataset, an image dataset, and a text dataset, this functionality would offer a clean and intuitive solution to create multimodal datasets. I am aware that the aforementioned is also feasible without my proposed function, but I believe this offers a nice approach that aligns with existing functionality and is directly provided within the `datasets` package.

### API
`stack_datasets` has two arguments: `datasets` and `stopping_strategy `.
<br />

`datasets` is a dictionary of either type `Dict[str, Dataset]` or `Dict[str, IterableDatasets]`, a mixture is not allowed. It contains the names of the datasets (the keys) and the datasets themselves (the values) that should be stacked. Each item returned is a dictionary with one key-value pair for each dataset. The keys are the names of the datasets as provided in the argument `datasets`, and the values are the respective examples from the datasets.
<br />

`stopping_strategy` is the same as for `interleave_datasets`. If it is `first_exhausted` we stop if the smallest dataset runs out of examples, if it is `all_exhausted` we stop if all datasets ran out of examples at least once. For `all_exhausted` that means that we may visit examples from datasets multiple times. 

### Docs
I saw that there are multiple documentations and guides on the HuggingFace website that introduce `concatenate_datasets` and `interleave_datasets`, for example [here](https://huggingface.co/docs/datasets/process#concatenate). If this request is merged I would be willing to add the new functionality at the appropriate points in the documentation (if desired).

### Tests
I also added some tests to ensure correctness. Some tests I wrote in [tests/test_iterable_dataset.py](https://github.com/TimCares/datasets/blob/fadc1159debf2a65d44e40cbf7758f2bd2cc8b08/tests/test_iterable_dataset.py#L2169)
run for both `Dataset` and `IterableDataset` even though tests for `Dataset` technically do not belong in this script, but I found that this was a nice way to cover more cases with mostly the same code.

### Additional information
I tried to write the code in a way so that it is similar to that of `concatenate_datasets` and `interleave_datasets`.
I’m open to feedback and willing to make adjustments based on your suggestions, so feel free to give me your take. :)

Comments: