Dataset from file

[jlamypoirier]

✨ Description

Fix: #136

Replace the both the json and concatenated memmap dataset by a generic (fast-llm) config from file. The dataset preparator now generates such config file, a blended dataset with probabilities proportional to file sizes (same as json dataset). Note:

• Dropped the json format from the preparator, but existing json datasets still work.
• Deprecated the concatenated memmap format, keeping since @oleksost is using it but we should try to remove soon.
• Added a way to split the dataset in the preparator. This was necessary to allow train/valid/test (as is done in tests)
• Updated the tutorial to use the new format.

🔍 Type of change

Select all that apply:

• 🐛 Bug fix (non-breaking change that addresses a specific issue)
• 🚀 New feature (non-breaking change that adds functionality)
• ⚠️ Breaking change (a change that could affect existing functionality)
• 📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
• 🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
• 📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
• 📝 Documentation change (updates documentation, including new content or typo fixes)
• 🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

[tscholak]

Hi @jlamypoirier, thanks.
I'd like for @RaymondLi0 and @oleksost to have a look at these changes before we merge them.
Please also make sure the tests pass.

[tscholak]

this is the same code as above

[tscholak]

I dropped the concatenated memmap format, assuming no one is using it yet.

In fact @oleksost is using it, so we can't just remove it.

[jlamypoirier]

@tscholak @RaymondLi0 @oleksost PR is ready, can you please review?

[oleksost]

Regarding deprecating concatenated_memmap, I had a question here: #156 (comment)

[jlamypoirier]

Which dataset wrappers are currently available? Here we only document blended, file and sampled. Is the concatenated_memmap (will be deprecated) and memmap still available? Would it make sense to document each of the available wrappers?

They are still available but not expected to be used directly by normal users. memmap is should to be accessed through the yaml file, and concatenated_memmap is to be removed once we stop using it.

I think what's missing here is an example that shows how to turn a bunch of idx and bin files into a yaml or json dataset config file.

The yaml file is not to be created manually, it is created in dataset preparation so is already covered. Remains the case of already prepared datasets that don't have a yaml file, we can either re-prepare those or make a quick script that create it for an already prepared dataset, what do you think?

If I have a directory with a bunch of shards (shard_*.bin and shard_*.idx files) and the corresponding fast_llm_dataset.json (<- generated by the prepare command) file in it, is it equivalent to use:

type: file
path: path/to/directory/fast_llm_dataset.json

and

type: memmap
path: path/to/directory/shard

and

type: concatenated_memmap
path: path/to/directory/

This are all different. file uses the yaml file created in the prepare step, it won't work with older json files. Both file and legacy (json) are a probabilistic sampling of all shard.memmap is only a single shard (and not to be used in general), and concatenated_memmap is a single dataset made of all shards (and deprecated).

[jlamypoirier]

@tscholak @oleksost can we merge this?

[tscholak]

LGTM