-- extends "datasets/dataset_type.yaml"

[config_metadata]
    == super()
    -- set ns.config_name = "Small LM Interleaved"
    -- set ns.config_description = "Interleave all Small-LM datasets"
    -- set ns.main_feature = "text"

[main_body]
.define: &tokenizer !var "tokenizer"

    [pp_kwargs]
.define: &pp_kwargs !dict

    [fineweb]
.define: &fineweb !call:forgather:from_project
    <<: &dataset_args
        preprocess_args: !var "preprocess_args"
        tokenizer: *tokenizer
        pp_kwargs: *pp_kwargs
        shard_dataset: !var [ "shard_dataset", False ]
        select_range: !var [ "select_range", null ]
        shuffle: !var [ "shuffle", False ]
        seed: !var [ "seed", 42 ]
    project_dir: "{{ project_dir }}"
    config_template: "smollm-corpus/fineweb-edu-dedup.yaml"
    targets: [ "train_dataset", "eval_dataset", "test_dataset" ]

    [cosmopedia]
.define: &cosmopedia !call:forgather:from_project
    <<: *dataset_args
    project_dir: "{{ project_dir }}"
    config_template: "smollm-corpus/cosmopedia-v2.yaml"
    targets: [ "train_dataset", "eval_dataset", "test_dataset" ]

    [train_dataset]
train_dataset: &train_dataset !singleton:forgather.ml.datasets:interleave_datasets
    <<: &interleave_args
        # Probabilities are proportional to remaining examples in each dataset.
        probabilities: !partial:forgather.ml.datasets:balance_remaining_examples
        seed: 42
        stopping_strategy: "all_exhausted"
    datasets:
        - !call:getitem [ *fineweb, 'train_dataset' ]
        - !call:getitem [ *cosmopedia, 'train_dataset' ]

    [eval_dataset]
eval_dataset: !singleton:forgather.ml.datasets:interleave_datasets
    <<: *interleave_args
    datasets:
        - !call:getitem [ *fineweb, 'eval_dataset' ]
        - !call:getitem [ *cosmopedia, 'eval_dataset' ]

    [test_dataset]
test_dataset: !singleton:forgather.ml.datasets:interleave_datasets
    <<: *interleave_args
    datasets:
        - !call:getitem [ *fineweb, 'test_dataset' ]
        - !call:getitem [ *cosmopedia, 'test_dataset' ]