Multi-`train_dataset` in Huggingface Trainer
Author: sbmaruf
Pull request: no

# 🚀 Feature request

The current Huggingface Trainer Supports, a single `train_dataset` (torch.utils.data.dataset.Dataset). While it makes sense for most of the training setups, there are still some cases where it is convenient to have a list of `train_dataset`. The trainer can randomly select or follow a specific sampling strategy to select the samples from each of the `train_dataset`. Usually the papers mentioned in the `motivation` sections use a multinomial distribution with a penalty hyperparameter (\alpha). An example is attached below with code.

## Motivation

1. Easy Multi-task learning setup.
2. Multi-lingual pre-training mentioned in [XLM](https://github.com/facebookresearch/XLM), [mT5](https://arxiv.org/abs/2010.11934)
3. Even for LM fine-tuning [MultiMix](https://arxiv.org/abs/2004.13240) requires this feature.


## Your contribution

The sampling strategy for each of the  `train_dataset` (torch.utils.data.dataset.Dataset) can be varied by a penalty variable (\alpha). The sample code for multinomial distribution based sampling strategy is below,

```
def multinomial_prob(dataset_len, alpha=.5):
    tot_number_of_sent_in_all_lang = 0
    prob = OrderedDict()
    for k, v in dataset_len.items():
        tot_number_of_sent_in_all_lang += v
    for k, v in dataset_len.items():
        neu = v
        den = tot_number_of_sent_in_all_lang
        p = neu/den
        prob[k] = p

    q = OrderedDict()
    q_den = 0.0
    for k, v in prob.items():
        q_den += (v**alpha)
    sum_ = 0.0
    for k, v in prob.items():
        q[k] =  (v**alpha)/q_den
        sum_ += q[k]
    assert math.fabs(1-sum_) < 1e-5
    return q

```

```
def iterator_selection_prob(alpha, train_datasets, logger=None):
    dataset_len = OrderedDict()
    for k, v in train_datasets.items():
        dataset_len[k] = len(v)
    for k, v in dataset_len.items():
        logger.info("Total Number of samples in {} : {}".format(k, v))
    prob = multinomial_prob(dataset_len, alpha=alpha)
    logger.info("Language iterator selection probability.")
    ret_prob_index, ret_prob_list  = [], []
    for k,v in prob.items():
        ret_prob_index.append(k)
        ret_prob_list.append(v)
    for k, v in zip(ret_prob_index, ret_prob_list):
        logger.info("{} : {}".format(k, v))
    return dataset_len, ret_prob_index, ret_prob_list
```

Inside the training loop, we could integrate like the following (the sample code may not match with the `Trainer` code). This is just an example. 

```
for step in range(args.max_steps*args.gradient_accumulation_steps):
        model.train()
        iterator_id = np.random.choice(range(tot_num_of_iterator), p=lang_prob)
        try:
            batch = train_iterators[iterator_id].__next__()
        except StopIteration:
            train_iterators[iterator_id] = iter(train_data_loader[iterator_id][1])
            batch = train_iterators[iterator_id].__next__()
        num_of_batch_trained[ iterator_id ] += 1
```

Comments:
This can all be done in one `Dataset` that randomly picks elements from subdatasets, so there is no need to add anything to the `Trainer` to support this.

Hi! Sorry to bother you again. I could not find any example code, that's why I opened the issue.
Later after your comment, I search through the [repo](https://huggingface.co/docs/datasets/) but could not find any class name SubDataset [here](https://huggingface.co/docs/datasets/search.html?q=SubDataset&check_keywords=yes&area=default#).

After Searching in the repository, I found some related [examples](https://github.com/huggingface/datasets/blob/67574a8d74796bc065a8b9b49ec02f7b1200c172/datasets/wmt16/wmt_utils.py) with the same `SubDataset` keyword. Is that what you mean? 

@sgugger 

All the links point to the Datasets library, so you shuold maybe open an issue there or ask on the Datasets category of the [forums](https://discuss.huggingface.co/)?

Sure, thank you for the reply and closing the issue.