"""Module for constants related to the FineMath dataset."""

from mt_dao.dataset.constants.dataset_constants_types import (
    TRAIN_CONSTANT,
    DatasetConstants,
    DataSplitConstants,
)

DATASET_PATH = "HuggingFaceTB/finemath"
FINEMATH_3PLUS = "finemath-3plus"
FINEMATH_4PLUS = "finemath-4plus"
INFIWEBMATH_3PLUS = "infiwebmath-3plus"
INFIWEBMATH_4PLUS = "infiwebmath-4plus"


finemath_finemath3plus_constants = DatasetConstants(
    splits={
        TRAIN_CONSTANT: DataSplitConstants(
            path=DATASET_PATH,
            name=FINEMATH_3PLUS,
            split=TRAIN_CONSTANT,
            folder_split=TRAIN_CONSTANT,
            truncated_samples=None,
        ),
    },
)


finemath_finemath4plus_constants = DatasetConstants(
    splits={
        TRAIN_CONSTANT: DataSplitConstants(
            path=DATASET_PATH,
            name=FINEMATH_4PLUS,
            split=TRAIN_CONSTANT,
            folder_split=TRAIN_CONSTANT,
            truncated_samples=None,
        ),
    },
)

finemath_infiwebmath3plus_constants = DatasetConstants(
    splits={
        TRAIN_CONSTANT: DataSplitConstants(
            path=DATASET_PATH,
            name=INFIWEBMATH_3PLUS,
            split=TRAIN_CONSTANT,
            folder_split=TRAIN_CONSTANT,
            truncated_samples=None,
        ),
    },
)


finemath_infiwebmath4plus_constants = DatasetConstants(
    splits={
        TRAIN_CONSTANT: DataSplitConstants(
            path=DATASET_PATH,
            name=INFIWEBMATH_4PLUS,
            split=TRAIN_CONSTANT,
            folder_split=TRAIN_CONSTANT,
            truncated_samples=None,
        ),
    },
)
