"""Module to store constants for the datasets."""

from mt_dao.dataset.constants.dataset_constants_types import DatasetConstants
from mt_dao.dataset.constants.finemath import (
    finemath_finemath3plus_constants,
    finemath_finemath4plus_constants,
    finemath_infiwebmath3plus_constants,
    finemath_infiwebmath4plus_constants,
)
from mt_dao.dataset.constants.mc4 import (
    c4_de_constants,
    c4_el_constants,
    c4_en_constants,
    c4_es_constants,
    c4_hi_constants,
    c4_it_constants,
    c4_la_constants,
    c4_ms_constants,
    c4_ru_constants,
    c4_sr_constants,
    c4_sw_constants,
    c4_ur_constants,
    c4_zh_constants,
)
from mt_dao.dataset.constants.smollm_corpus import (
    smollm_corpus_cosmo_constants,
    smollm_corpus_fineweb_edu_dedup_constants,
    smollm_corpus_python_edu_constants,
)
from mt_dao.dataset.constants.the_pile import pile_constants

DATASETS_CONSTANTS: dict[str, DatasetConstants] = {
    "c4_en": c4_en_constants,
    "c4_it": c4_it_constants,
    "c4_zh": c4_zh_constants,
    "c4_ms": c4_ms_constants,
    "c4_ur": c4_ur_constants,
    "c4_sw": c4_sw_constants,
    "c4_la": c4_la_constants,
    "c4_sr": c4_sr_constants,
    "c4_es": c4_es_constants,
    "c4_de": c4_de_constants,
    "c4_el": c4_el_constants,
    "c4_ru": c4_ru_constants,
    "c4_hi": c4_hi_constants,
    "the_pile": pile_constants,
    "smollm_cosmo": smollm_corpus_cosmo_constants,
    "smollm_python_edu": smollm_corpus_python_edu_constants,
    "smollm_fineweb_edu_dedup": smollm_corpus_fineweb_edu_dedup_constants,
    "fine_math_fine_math_3plus": finemath_finemath3plus_constants,
    "fine_math_fine_math_4plus": finemath_finemath4plus_constants,
    "fine_math_infiwebmath_3plus": finemath_infiwebmath3plus_constants,
    "fine_math_infiwebmath_4plus": finemath_infiwebmath4plus_constants,
}
