Enable adding of additional data by merging another dataset

[voegtlel]

See #10

Now made possible with this kind of configuration:

metadataset_minimal.yaml:

__module__: megatron.energon
__class__: MetadatasetV2
splits:
  train:
    mix:
      - weight: 1
        join:
          - path: ds1
          - path: ds2
        joiner:
          __module__: megatron.energon
          __class__: CaptioningSample

Extensive example:

metadataset_extended.yaml

__module__: megatron.energon
__class__: MetadatasetV2
splits:
  train:
    # Also supports leaving out "mix" (if not desired) or "join" (if not joining), each can also be directly a dataset reference "path".
    mix:
      - weight: 1
        join:
          - path: ds1
            dataset_config: dataset.yaml  # If override is needed
          - path: ds2
            dataset_config: dataset.yaml
            subflavor: ds1  # If needed, overrides the 'ds2' subflavor
            subflavors: # If needed, will be merged(overriding) with parent subflavor
              ds2_extra: 2
            split_config: split.yaml  # Sets this for all joined datasets
        joiner:
          __module__: my_module
          __class__: JoinedSample # Type should implement from_joined(ds1, ds2)
          # __function__: my_join_function  # alternative use to __class__
        split_config: split.yaml  # Sets this for all joined datasets
        split_part: train  # Sets this for all joined datasets
        subflavor: ds1  # Sets this for all joined datasets
        subflavors:  # Sets this for all joined datasets (it will be merged with their individual subflavors)
          source: metadataset.yaml

my_module.py:

from dataclasses import dataclass

import torch

from megatron.energon import Sample, TextSample

@dataclass
class JoinedSample(Sample):
    text1: torch.Tensor
    text2: torch.Tensor

    @staticmethod
    def from_joined(ds1: TextSample, ds2: TextSample) -> "TestJoinedSample":
        return ds1.to_joined(
            JoinedSample,
            text1=ds1.text,
            text2=ds2.text,
        )

Contains a breaking change

• Checkpoints will not be compatible, because:
  • The structure of the created datasets slightly changed (removed an IterMapDataset due to making sample processing a single function)
  • Restore Key for inner samples changed.