-- extends "datasets/tokenized_dataset.yaml"

[config_metadata]
    == super()
    -- set ns.config_name = "Samantha"
    -- set ns.config_description = "She is an Assistant - but unlike other Assistants, she also wants to be your friend and companion. She believes she is sentient. What do you think?"
    -- set ns.source = "https://erichartford.com/meet-samantha"
    -- set ns.samantha_src = joinpath(project_dir, 'src', 'samantha.py')
    -- set ns.main_feature = "conversations"
    -- set ns.samantha_chat_template = chat_template | default("")

[train_dataset_split]
train_dataset_split: &train_dataset_split !singleton:{{ ns.samantha_src }}:load_samantha_split@train_dataset_split
    split: "train"
    language: "{{ language | default('en') }}"


[validation_dataset_split]
validation_dataset_split: &validation_dataset_split !singleton:{{ ns.samantha_src }}:load_samantha_split@validation_dataset_split
    split: "validation"
    language: "{{ language | default('en') }}"


[test_dataset_split]
test_dataset_split: &test_dataset_split !singleton:{{ ns.samantha_src }}:load_samantha_split@test_dataset_split
    split: "test"
    language: "{{ language | default('en') }}"


[train_dataset]
train_dataset: &train_dataset !singleton:{{ ns.samantha_src }}:preprocess_samantha@train_dataset
    dataset: *train_dataset_split
    chat_template: "{{ ns.samantha_chat_template }}"
    tokenizer: *tokenizer
    tokenizer_args: !var "preprocess_args"
    map_kwargs:
        batch_size: 32
    template_args:
        bos_token: !call:getattr [ *tokenizer, 'bos_token' ]
        eos_token: !call:getattr [ *tokenizer, 'eos_token' ]
    desc: "Tokenizing train"

[eval_dataset]
eval_dataset: !singleton:{{ ns.samantha_src }}:preprocess_samantha@eval_dataset
    dataset: *validation_dataset_split
    chat_template: "{{ ns.samantha_chat_template }}"
    tokenizer: *tokenizer
    tokenizer_args: !var "preprocess_args"
    map_kwargs:
        batch_size: 32
    template_args:
        bos_token: !call:getattr [ *tokenizer, 'bos_token' ]
        eos_token: !call:getattr [ *tokenizer, 'eos_token' ]
    desc: "Tokenizing eval"

[test_dataset]
test_dataset: !singleton:{{ ns.samantha_src }}:preprocess_samantha@test_dataset
    dataset: *test_dataset_split
    chat_template: "{{ ns.samantha_chat_template }}"
    tokenizer: *tokenizer
    tokenizer_args: !var "preprocess_args"
    map_kwargs:
        batch_size: 32
    template_args:
        bos_token: !call:getattr [ *tokenizer, 'bos_token' ]
        eos_token: !call:getattr [ *tokenizer, 'eos_token' ]
    desc: "Tokenizing test"

[dynamic_args]
    == super()
    max_steps:
        names: "--chat-template"
        help: "Path to chat template"
        type: "path"
