Multi-Task dataset with Custom Sampler and Sharding

The current Huggingface Trainer Supports, a single train_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. An example is attached below for the custom sampling streategy code.

The sampling strategy for each of the train_dataset ( (from multiple train dataset) can be varied by a penalty variable (\alpha). The sample code for a custom multinomial distribution based sampling strategy is below (mentioned in XLM paper),

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():"Total Number of samples in {} : {}".format(k, v))
    prob = multinomial_prob(dataset_len, alpha=alpha)"Language iterator selection probability.")
    ret_prob_index, ret_prob_list  = [], []
    for k,v in prob.items():
    for k, v in zip(ret_prob_index, ret_prob_list):"{} : {}".format(k, v))
    return dataset_len, ret_prob_index, ret_prob_list

So I have three questions in general,

  1. How to integrate multiple datasets (or sub-dataset) in the same dataset class?
  2. How to apply custom control on the sampling strategy (let’s just say I want to inject the above sampling strategy in my sub-datasets) into different sub-dataset?
  3. Also in the case of the large tenderized dataset that cannot fit into memory how to handle the sharding using huggingface trainer.


  1. I am not looking for sample codes. A discussion or pointer to the Hf source library is also highly appreciated. However, sample codes are always best.
  2. I would also like to know if you have seen some other repository implements these feature with/without Hf-library.
  3. Discussion on any topic is highly appreciated.

@sgugger Do you have any idea on these topics?