Add sample method to data_pipeline

[najielhachem]

What does this PR do? Please describe:
Add sample data source operator that can sample from multiple data sources provided weights.

Does your PR introduce any breaking changes? If yes, please list them:
List of all backwards-incompatible changes.

Check list:

• Was the content of this PR discussed and approved via a GitHub issue? (no need for typos or documentation improvements)
• Did you read the contributor guideline?
• Did you make sure that your PR does only one thing instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests?
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (no need for typos, documentation, or minor internal changes)

[gwenzek]

The code LGTM, but I wanted to talk about the use case and the handling of end of pipeline.

In this implementation, as soon as one child data source is done, the sampled_data_source stops.
In round_robin_data_source on the contrary I ensure that all data source have been exhausted at least once before stopping.

• do we need the two logics ?
• if yes, I think they should be in the same data source with a stop_on_shortest flag
• if yes, I think we could merge sampled_data_source and round_robin_data_source by adding a weights parameter to round_robin

[gwenzek]

could you also document the behavior around the end of the pipeline ?
when is this pipeline considered exhausted ?

[cbalioglu]

@gwenzek although they are technically similar in that they both iterate through various data pipelines in each iteration, I think semantically they are fundamentally two different operations. For sampling you never want to keep reading from a data pipeline once it is exhausted, doing that would defeat the whole purpose of sampling. So to me stop_on_shortest feels like a redundant/confusing parameter for sampling.

Maybe they can internally share some implementation logic, but I think they should be exposed as two separate operators.

[cbalioglu]

Thanks again for your work Naji! And sorry for the late review. I just left an important comment that should be easy to fix, otherwise looks good to me.

[gwenzek]

For sampling you never want to keep reading from a data pipeline once it is exhausted, doing that would defeat the whole purpose of sampling

imagine you're training a multilingual model. You have English, French and Zulu data. Volume is 100Gb of English, 10Gb of French, 1Gb of Zulu.
You want to up sample the Zulu and down sample English. So that you're model is trained on 50% English, 40% French and 10% Zulu. If you use sample([0.5, 0.4, 0.1]) you're epoch will stop after having read the first 5Gb of English, which is only 5% of the original training set. If you then reset the dataloader and start a new epoch you'll restart reading the first 5Gb of English over and over.

I think semantically they are fundamentally two different operations

Round robin looks like a lot like .sample([1/n, 1/n, ... 1/n]) minus the randomization part

[cbalioglu]

You want to up sample the Zulu and down sample English. So that you're model is trained on 50% English, 40% French and 10% Zulu. If you use sample([0.5, 0.4, 0.1]) you're epoch will stop after having read the first 5Gb of English, which is only 5% of the original training set. If you then reset the dataloader and start a new epoch you'll restart reading the first 5Gb of English over and over.

Yeah that makes sense. I haven't been focusing on up/down sampling really, but when you consider that, stop_on_shortest sounds reasonable. Thanks for the example.

I still think that round-robin and sampling should be two separate ops though. Round-robin by its nature does not really have any randomness, it is all about iterating each pipeline one-by-one. It feels weird to me to use an op called round_robin() for sampling or the other way round.

[najielhachem]

I agree with @cbalioglu about having two operators since semantically we don't expect a similar behaviour from round_robin and sample. But as mentioned by @gwenzek, the buffering seems useful so I will extract this part into a separate code and use it in both of our operators.

[cbalioglu]

LGTM!