Add metadata to dataset config files

[jlamypoirier]

✨ Description

Gather all available metadata from readers during dataset preparation and add it to the yaml config file. Works with both blending and splitting. Example:

$ cat /tmp/fast_llm_tests/common/dataset/dataset_with_image_patches_55_132/fast_llm_config.yaml 
config:
  path: shard_0_0.fast_llm_dataset
  type: memmap
metadata:
  image_patches:
    data_type: uint8
    num_documents: 1000
    num_patch_groups: 1355
    num_patches: 10707
    num_pixels: 513936
    patch_shape:
    - 3
    - 4
    - 4
  num_tokens: 59145
  tokens:
    data_type: int32
    num_documents: 1000
    num_tokens: 59145

The file format is different from before (config is now inside of a dict instead of at top level), but the old format can still be loaded.

[tscholak]

This looks very useful, thanks @jlamypoirier
@oleksost @RaymondLi0 is this close to what you had in mind?

[oleksost]

yes, this looks useful.
So the fast_llm_config.yaml will contain such entry per shard? Wouldn't it be more useful to have a dataset view isntead of per shard view?

[jlamypoirier]

yes, this looks useful. So the fast_llm_config.yaml will contain such entry per shard? Wouldn't it be more useful to have a dataset view isntead of per shard view?

This is already global, the example just happens to have only one shard.

[oleksost]

There are some merge conflicts, apart from that LHTM!

[tscholak]

@jlamypoirier this is causing crashes now. remove?

Suggested change

Assert.eq(config.keys(), {"config", "metadata"})

[jlamypoirier]

Yes, I must have left this there by accident