Arbitrary Dataset Mixtures with DataModules

[siddk]

Hey folks - I have a semi-standard usecase that I was hoping to use DataModules for. I'd like to train a language model on multiple different training datasets (each processed individually), sampling my full batch from each sub-dataset according to some mixture parameters (e.g. 50% from dataset 1, 25% from dataset 2, 25% from dataset 3).

Is there a nice way to do this with DataModules? Separately, if I wanted to extend this to streaming datasets, what would I need to do?

[rohitgr7]

do you mean the distribution should be used across training or per epoch?

like choose dataset1 for 50% of all epochs, 25% for dataset2....
or in each epoch, choose dataset1 50% of total when a sample is requested, 25% dataset2....

[siddk]

I was thinking the latter -- throughout training each batch is comprised of 50% dataset1, 25% dataset2...

The usecase is that a lot of language modeling datasets come from sources of various quality; I'd want to upsample the high-quality sources within each batch, but still see the full diversity of data.