Source code for kempnerforge.config.data

"""Data pipeline configuration."""

from __future__ import annotations

from dataclasses import dataclass, field

TOKEN_DTYPES = ("uint16", "uint32")




def check_token_dtype(value: str, key: str) -> None:
    """Reject a ``token_dtype`` that is neither ``""`` nor one of ``TOKEN_DTYPES``."""
    if value != "" and value not in TOKEN_DTYPES:
        raise ValueError(
            f"{key} must be 'uint16' or 'uint32', the integer type the raw .bin token "
            f"shards were written with, or empty to take it from the metadata.yaml next "
            f"to the shards; got {value!r}"
        )





@dataclass
class DatasetSource:
    """A single data source in a multi-dataset mixture.

    Either ``path`` (pre-tokenized) or ``hf_name`` (HuggingFace) must be set.
    ``weight`` controls the relative sampling probability. ``token_dtype``
    overrides ``data.token_dtype`` for the ``.bin`` shards under ``path``.
    """

    path: str = ""
    weight: float = 1.0
    name: str = ""
    hf_name: str = ""
    hf_config: str = ""
    token_dtype: str = ""

    def __post_init__(self) -> None:
        if self.weight <= 0:
            raise ValueError(f"Dataset weight must be positive, got {self.weight}")
        check_token_dtype(self.token_dtype, "[[data.datasets]] token_dtype")





@dataclass
class TrainingPhase:
    """A training phase with custom dataset weights and LR scaling.

    Used for data annealing: at ``start_step``, the mixture sampler switches
    to ``dataset_weights`` and the learning rate is multiplied by ``lr_scale``.
    """

    start_step: int = 0
    dataset_weights: dict[str, float] = field(default_factory=dict)
    lr_scale: float = 1.0

    def __post_init__(self) -> None:
        if self.start_step < 0:
            raise ValueError("TrainingPhase.start_step must be non-negative")
        if self.lr_scale <= 0:
            raise ValueError("TrainingPhase.lr_scale must be positive")
        for name, w in self.dataset_weights.items():
            if w < 0:
                raise ValueError(f"TrainingPhase.dataset_weights['{name}'] must be non-negative")





@dataclass
class DataConfig:
    """Data pipeline settings.

    ``token_dtype`` is the integer type of the raw ``.bin`` token shards,
    ``"uint16"`` or ``"uint32"``, for ``dataset_path`` and every ``datasets``
    source that does not set its own. Left empty, the ``metadata.yaml`` next to
    the shards decides, as ``MemoryMappedDataset`` describes. ``.npy`` shards
    carry their dtype in their header and ignore it.
    """

    dataset_path: str = ""
    file_pattern: str = "*.npy"
    token_dtype: str = ""
    tokenizer_path: str = ""
    num_workers: int = 4
    pin_memory: bool = True
    prefetch_factor: int = 2
    hf_dataset_name: str | None = None
    hf_dataset_config: str | None = None
    hf_dataset_split: str = "train"
    hf_dataset_text_field: str = "text"
    hf_dataset_image_field: str = "image"
    hf_dataset_prompt_field: str = ""
    hf_image_size: int = 224
    hf_streaming: bool = False
    pack_sequences: bool = False
    datasets: list[DatasetSource] = field(default_factory=list)
    mix_temperature: float = 1.0
    phases: list[TrainingPhase] = field(default_factory=list)
    anneal_start_step: int = 0
    anneal_weights: dict[str, float] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.num_workers < 0:
            raise ValueError("num_workers must be non-negative")
        if self.prefetch_factor < 1:
            raise ValueError("prefetch_factor must be >= 1")
        if self.mix_temperature <= 0:
            raise ValueError("mix_temperature must be positive")
        check_token_dtype(self.token_dtype, "data.token_dtype")
        for src in self.datasets:
            if not src.path and not src.hf_name:
                raise ValueError(f"DatasetSource '{src.name}' must have either path or hf_name")
        if self.phases:
            steps = [p.start_step for p in self.phases]
            if steps != sorted(steps) or len(steps) != len(set(steps)):
                raise ValueError(
                    "data.phases start_steps must be strictly monotonically increasing"
                )
        if self.anneal_start_step < 0:
            raise ValueError("anneal_start_step must be non-negative")
        if self.phases and self.anneal_start_step > 0:
            raise ValueError(
                "Cannot use both data.phases and data.anneal_start_step — "
                "use phases for multi-phase scheduling"
            )