dataset blending tool

[RaymondLi0]

✨ Description

Tool to recursively discover datasets in a directory and generate a blended dataset config.
This tool walks through a directory tree, identifies datasets by their fast_llm_config*.yaml files,
and generates a config file that blends all discovered datasets with weights proportional to token counts.

🔍 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]

great!
though why isn't that built into fast-llm itself? ideally we shouldn't need this.
cc @jlamypoirier

[jlamypoirier]

great! though why isn't that built into fast-llm itself? ideally we shouldn't need this. cc @jlamypoirier

@tscholak @RaymondLi0 we already have all the machinery we need for blending datasets in GPTMemmapDatasetPreparator, anb ideally we should extract and reuse it. What is missing for a standalone blending tool is dataset discovery and reading.

The discovery implementation here looks at the yaml configs which cause complications because of all the ways they could be defined. A much simpler option would be to just look at the .fast_llm_dataset files and blend them.

As for reading, it should be done by instantiating the dataset, it will take care of everything and avoid trouble.

Lastly, note that almost every "tool" we've had so far quickly went out of sync and became useless. To prevent this script from suffering the same fate, we'll want to bring it inside fast_llm (would be appropriate as a DatasetPreparator) and add tests for it..

[RaymondLi0]

Made the following updates:

• moved the code under fast_llm/data/preparator as a DatasetPreparator subclass.
• find .fast_llm_dataset files instead of yaml configs
• add tests

[jlamypoirier]

Not needed, you should be able to instantiate the dataset and get its num_tokens

[RaymondLi0]

Actually the compatibility check here: https://github.com/ServiceNow/Fast-LLM/blob/main/fast_llm/data/preprocessing/language_model.py#L35
prevents to instantiate the dataset without reading the config from the .fast_llm_dataset file it seems

[RaymondLi0]

Refactored it a bit to remove some redundant code

[jlamypoirier]

Why making this hierarchical? It's not really needed and it's making this class way more complicated than it needs to be, plus nested blended datasets aren't really a good idea and comes with important downsides like extra overhead and randomness issues.

[RaymondLi0]

By creating groups of datasets, the hierarchical part makes it much more convenient to update the mixing weights.

[RaymondLi0]

Outputting a flat config would defeat the purpose of this tool, which is to enable easy configuration and manual update of the dataset mix.

[jlamypoirier]

Better leave as integer as in gpt memmap preparator.

[RaymondLi0]

This is also to facilitate human review of the configs. Is there some reason to keep it as an integer, other than consistency with gpt-memmap?

[jlamypoirier]

Main goal was to avoid rounding errors and make the token count explicit, when compared to normalizing probabilities. The division by 1e9 doesn't prevent this though so is probably fine.