search

Lightning-AI / litdata

Streamline data pipelines for AI. Process datasets across 1000s of machines, and optimize data for blazing fast model training.
Apache License 2.0
250 stars 24 forks source link

Add support for `iterate_over_all` for the CombinedDataset #122

Closed tchaton closed 1 month ago

tchaton commented 1 month ago
Before submitting - [ ] Was this discussed/agreed via a Github issue? (no need for typos and docs improvements) - [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/lit-data/blob/main/.github/CONTRIBUTING.md), Pull Request section? - [ ] Did you make sure to update the docs? - [ ] Did you write any new necessary tests?

This PR aims at improving the behaviour of the combined dataset for more traditional use case.

Before this PR, the CombinedDataset would stop iterating as soon as one of its wrapped dataset would trigger a stop iteration. This is a great for LLM pre-training but less practical for more traditional use case.

This PR introduces the argument iterate_over_all to enable the combined dataset to see all the samples from all the datasets. The sampling is still random but only the last items will be un-sampled.

What does this PR do?

Fixes #112

PR review

Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃