Blame for runtime/backends/actioncodec2/configuration.py

"""Validated ``config.yaml`` model contract for ActionCodec2."""

from __future__ import annotations

import math
import re
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType

import yaml

from ._validation import validate_dimension_partition

CONFIG_NAME = "config.yaml"

MODEL_TYPE = "actioncodec2_unified"
CONFIG_VERSION = 1
_SPACES = ("additive", "so3")
_SO3_REPRESENTATIONS = ("rotation_vector", "euler_xyz")
_SO3_MULTIPLICATIONS = ("left", "right")
_PRIMITIVE_SELECTIONS = ("nearest", "smooth", "momentum_lazy")
_ORDER_KEY = re.compile(r"order_(0|[1-9][0-9]*)\Z")

NumberOrVector = float | tuple[float, ...]


class ConfigError(ValueError):
    """The serialized model definition is incomplete or inconsistent."""


class _UniqueKeyLoader(yaml.SafeLoader):
    pass


def _construct_unique_mapping(
    loader: _UniqueKeyLoader, node: yaml.MappingNode, deep: bool = False
) -> dict[object, object]:
    loader.flatten_mapping(node)
    output: dict[object, object] = {}
    for key_node, value_node in node.value:
        key = loader.construct_object(key_node, deep=deep)
        try:
            duplicate = key in output
        except TypeError as error:
            raise ConfigError("config.yaml contains an invalid mapping key") from error
        if duplicate:
            raise ConfigError(f"config.yaml contains duplicate field {key!r}")
        output[key] = loader.construct_object(value_node, deep=deep)
    return output


_UniqueKeyLoader.add_constructor(
    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
    _construct_unique_mapping,
)


def _mapping(value: object, path: str) -> Mapping[str, object]:
    if not isinstance(value, Mapping):
        raise ConfigError(f"{path} must be a mapping")
    if not all(isinstance(key, str) for key in value):
        raise ConfigError(f"every field in {path} must be a string")
    return value


def _reject_unknown(payload: Mapping[str, object], allowed: set[str], path: str) -> None:
    unknown = sorted(set(payload) - allowed)
    if unknown:
        rendered = ", ".join(unknown)
        raise ConfigError(f"{path} contains unknown field(s): {rendered}")


def _required(payload: Mapping[str, object], key: str, path: str) -> object:
    if key not in payload:
        raise ConfigError(f"{path} is missing required field {key!r}")
    return payload[key]


def _integer(value: object, path: str, *, positive: bool = False) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise ConfigError(f"{path} must be an integer")
    minimum = 1 if positive else 0
    if value < minimum:
        qualifier = "positive" if positive else "non-negative"
        raise ConfigError(f"{path} must be {qualifier}")
    return int(value)


def _number(value: object, path: str) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ConfigError(f"{path} must be a number")
    result = float(value)
    if not math.isfinite(result):
        raise ConfigError(f"{path} must be finite")
    return result


def _number_or_vector(value: object, width: int, path: str) -> NumberOrVector:
    if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
        result = tuple(_number(item, f"{path}[{index}]") for index, item in enumerate(value))
        if len(result) != width:
            raise ConfigError(f"{path} must be a scalar or contain {width} values")
        return result
    return _number(value, path)


def _expanded(value: NumberOrVector, width: int) -> tuple[float, ...]:
    if isinstance(value, tuple):
        return value
    return (value,) * width


def _serialize_number_or_vector(value: NumberOrVector) -> float | list[float]:
    return list(value) if isinstance(value, tuple) else value


def _parse_dims(value: object, input_dim: int, path: str) -> tuple[int, ...]:
    if isinstance(value, Mapping):
        payload = _mapping(value, path)
        _reject_unknown(payload, {"start", "stop"}, path)
        start = _integer(_required(payload, "start", path), f"{path}.start")
        stop = _integer(_required(payload, "stop", path), f"{path}.stop")
        if stop <= start:
            raise ConfigError(f"{path}.stop must be greater than {path}.start")
        dims = tuple(range(start, stop))
    elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
        dims = tuple(_integer(item, f"{path}[{index}]") for index, item in enumerate(value))
        if not dims:
            raise ConfigError(f"{path} must not be empty")
    else:
        raise ConfigError(f"{path} must be a list or a start/stop range")

    if len(set(dims)) != len(dims):
        raise ConfigError(f"{path} contains a repeated dimension")
    outside = [dimension for dimension in dims if dimension >= input_dim]
    if outside:
        raise ConfigError(f"{path} contains dimension {outside[0]}, outside input_dim={input_dim}")
    return dims


def _serialize_dims(dims: tuple[int, ...]) -> object:
    if len(dims) > 1 and dims == tuple(range(dims[0], dims[-1] + 1)):
        return {"start": dims[0], "stop": dims[-1] + 1}
    return list(dims)


@dataclass(frozen=True)
class ConstraintConfig:
    """Allowed inclusive interval for one differential order."""

    minimum: NumberOrVector
    maximum: NumberOrVector

    @classmethod
    def from_dict(cls, payload: object, *, width: int, path: str) -> ConstraintConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"min", "max"}, path)
        minimum = _number_or_vector(_required(values, "min", path), width, f"{path}.min")
        maximum = _number_or_vector(_required(values, "max", path), width, f"{path}.max")
        if any(
            lower > upper
            for lower, upper in zip(
                _expanded(minimum, width), _expanded(maximum, width), strict=True
            )
        ):
            raise ConfigError(f"{path}.min must not exceed {path}.max")
        return cls(minimum=minimum, maximum=maximum)

    def to_dict(self) -> dict[str, object]:
        return {
            "min": _serialize_number_or_vector(self.minimum),
            "max": _serialize_number_or_vector(self.maximum),
        }


@dataclass(frozen=True)
class InputTransformConfig:
    """Input-unit conversion and optional source-side clipping."""

    scale: NumberOrVector = 1.0
    clip: bool = False

    @classmethod
    def from_dict(cls, payload: object | None, *, width: int, path: str) -> InputTransformConfig:
        if payload is None:
            return cls()
        values = _mapping(payload, path)
        _reject_unknown(values, {"scale", "clip"}, path)
        scale = _number_or_vector(values.get("scale", 1.0), width, f"{path}.scale")
        if any(value <= 0.0 for value in _expanded(scale, width)):
            raise ConfigError(f"{path}.scale must be positive")
        clip = values.get("clip", False)
        if not isinstance(clip, bool):
            raise ConfigError(f"{path}.clip must be true or false")
        return cls(scale=scale, clip=clip)

    def to_dict(self) -> dict[str, object]:
        return {
            "scale": _serialize_number_or_vector(self.scale),
            "clip": self.clip,
        }


def _parse_component_geometry(
    values: Mapping[str, object], *, input_dim: int, path: str
) -> tuple[tuple[int, ...], str, int, int, str | None, str | None]:
    dims = _parse_dims(_required(values, "dims", path), input_dim, f"{path}.dims")
    space = _required(values, "space", path)
    if space not in _SPACES:
        raise ConfigError(f"{path}.space must be one of {_SPACES}")
    input_order = _integer(_required(values, "input_order", path), f"{path}.input_order")
    primitive_order = _integer(
        _required(values, "primitive_order", path), f"{path}.primitive_order"
    )
    representation = values.get("representation")
    multiplication = values.get("multiplication")
    if space == "so3":
        if len(dims) != 3:
            raise ConfigError(f"{path} with space='so3' must contain exactly 3 dims")
        if input_order not in {0, 1} or primitive_order not in {1, 2}:
            raise ConfigError(
                f"{path} with space='so3' supports input_order 0 or 1 and primitive_order 1 or 2"
            )
        if representation not in _SO3_REPRESENTATIONS:
            raise ConfigError(f"{path}.representation must be one of {_SO3_REPRESENTATIONS}")
        if multiplication not in _SO3_MULTIPLICATIONS:
            raise ConfigError(f"{path}.multiplication must be one of {_SO3_MULTIPLICATIONS}")
    else:
        if input_order not in {0, 1, 2} or primitive_order not in {0, 1, 2}:
            raise ConfigError(
                f"{path} with space='additive' supports input_order and primitive_order 0, 1, or 2"
            )
        if representation is not None or multiplication is not None:
            raise ConfigError(
                f"{path}.representation and multiplication are only valid for space='so3'"
            )
    return dims, str(space), input_order, primitive_order, representation, multiplication


def _parse_component_constraints(
    values: Mapping[str, object], *, width: int, path: str
) -> dict[str, ConstraintConfig]:
    raw = _mapping(_required(values, "constraints", path), f"{path}.constraints")
    if not raw:
        raise ConfigError(f"{path}.constraints must not be empty")
    constraints: dict[str, ConstraintConfig] = {}
    for key, constraint in raw.items():
        if _ORDER_KEY.fullmatch(key) is None:
            raise ConfigError(
                f"{path}.constraints field {key!r} must be named order_0, order_1, ..."
            )
        constraints[key] = ConstraintConfig.from_dict(
            constraint, width=width, path=f"{path}.constraints.{key}"
        )
    return constraints


@dataclass(frozen=True)
class ComponentConfig:
    """One orthogonal physical component of the action vector."""

    dims: tuple[int, ...]
    space: str
    input_order: int
    primitive_order: int
    input_transform: InputTransformConfig
    constraints: Mapping[str, ConstraintConfig]
    distortion_limit: NumberOrVector
    primitive_selection: str
    representation: str | None = None
    multiplication: str | None = None

    @classmethod
    def from_dict(
        cls,
        payload: object,
        *,
        input_dim: int,
        path: str,
    ) -> ComponentConfig:
        values = _mapping(payload, path)
        allowed = {
            "dims",
            "space",
            "input_order",
            "primitive_order",
            "input_transform",
            "constraints",
            "distortion_limit",
            "primitive_selection",
            "representation",
            "multiplication",
        }
        _reject_unknown(values, allowed, path)
        dims, space, input_order, primitive_order, representation, multiplication = (
            _parse_component_geometry(values, input_dim=input_dim, path=path)
        )
        width = len(dims)
        constraints = _parse_component_constraints(values, width=width, path=path)

        input_transform = InputTransformConfig.from_dict(
            values.get("input_transform"),
            width=width,
            path=f"{path}.input_transform",
        )
        if input_transform.clip and f"order_{input_order}" not in constraints:
            raise ConfigError(
                f"{path}.input_transform.clip requires constraints.order_{input_order}"
            )

        distortion_limit = _number_or_vector(
            _required(values, "distortion_limit", path),
            width,
            f"{path}.distortion_limit",
        )
        if any(value <= 0.0 for value in _expanded(distortion_limit, width)):
            raise ConfigError(f"{path}.distortion_limit must be positive")
        if space == "so3" and isinstance(distortion_limit, tuple):
            raise ConfigError(
                f"{path}.distortion_limit must be one scalar geodesic angle for SO(3)"
            )

        primitive_selection = _required(values, "primitive_selection", path)
        if primitive_selection not in _PRIMITIVE_SELECTIONS:
            raise ConfigError(
                f"{path}.primitive_selection must be one of {', '.join(_PRIMITIVE_SELECTIONS)}"
            )
        if primitive_selection == "momentum_lazy" and (space != "so3" or primitive_order != 1):
            raise ConfigError(
                f"{path}.primitive_selection='momentum_lazy' requires first-order SO(3) primitives"
            )
        if (
            primitive_selection == "smooth"
            and space == "additive"
            and primitive_order not in {1, 2}
        ):
            raise ConfigError(
                f"{path}.primitive_selection='smooth' requires first- or second-order additive primitives"
            )
        if primitive_order == 2 and "order_2" not in constraints:
            raise ConfigError(f"{path} second-order primitives require constraints.order_2")
        return cls(
            dims=dims,
            space=space,
            input_order=input_order,
            primitive_order=primitive_order,
            input_transform=input_transform,
            constraints=MappingProxyType(constraints),
            distortion_limit=distortion_limit,
            primitive_selection=str(primitive_selection),
            representation=None if representation is None else str(representation),
            multiplication=None if multiplication is None else str(multiplication),
        )

    def to_dict(self) -> dict[str, object]:
        payload: dict[str, object] = {
            "dims": _serialize_dims(self.dims),
            "space": self.space,
            "input_order": self.input_order,
            "primitive_order": self.primitive_order,
            "input_transform": self.input_transform.to_dict(),
            "constraints": {
                key: constraint.to_dict() for key, constraint in self.constraints.items()
            },
            "distortion_limit": _serialize_number_or_vector(self.distortion_limit),
            "primitive_selection": self.primitive_selection,
        }
        if self.space == "so3":
            payload["representation"] = self.representation
            payload["multiplication"] = self.multiplication
        return payload


@dataclass(frozen=True)
class PhysicalConfig:
    """Complete, non-overlapping component partition of the input's last axis."""

    input_dim: int
    components: Mapping[str, ComponentConfig]

    @classmethod
    def from_dict(cls, payload: object, *, path: str = "physical") -> PhysicalConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"input_dim", "components"}, path)
        input_dim = _integer(
            _required(values, "input_dim", path), f"{path}.input_dim", positive=True
        )
        raw_components = _mapping(_required(values, "components", path), f"{path}.components")
        if not raw_components:
            raise ConfigError(f"{path}.components must not be empty")

        components: dict[str, ComponentConfig] = {}
        for name, component_payload in raw_components.items():
            if not name.strip():
                raise ConfigError(f"{path}.components contains an empty component name")
            component = ComponentConfig.from_dict(
                component_payload,
                input_dim=input_dim,
                path=f"{path}.components.{name}",
            )
            components[name] = component

        validate_dimension_partition(
            ((name, component.dims) for name, component in components.items()),
            input_dim,
            label="physical dimension",
            error_type=ConfigError,
        )
        return cls(input_dim=input_dim, components=MappingProxyType(components))

    def to_dict(self) -> dict[str, object]:
        return {
            "input_dim": self.input_dim,
            "components": {
                name: component.to_dict() for name, component in self.components.items()
            },
        }


TOKENIZATION_SCHEMES = ("set_bpe", "unified_set_bpe", "bpe", "sparse_set_bpe")


@dataclass(frozen=True)
class SetBPEConfig:
    """Runtime settings of the tokenization stage."""

    lookahead: int
    allow_cross_component_merge: bool = False
    scheme: str = "set_bpe"

    @classmethod
    def from_dict(cls, payload: object, *, path: str = "set_bpe") -> SetBPEConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"lookahead", "allow_cross_component_merge", "scheme"}, path)
        lookahead = _integer(
            _required(values, "lookahead", path), f"{path}.lookahead", positive=True
        )
        allow_cross = values.get("allow_cross_component_merge", False)
        if not isinstance(allow_cross, bool):
            raise ConfigError(f"{path}.allow_cross_component_merge must be true or false")
        scheme = values.get("scheme", "set_bpe")
        if scheme not in TOKENIZATION_SCHEMES:
            raise ConfigError(f"{path}.scheme must be one of {', '.join(TOKENIZATION_SCHEMES)}")
        return cls(lookahead=lookahead, allow_cross_component_merge=allow_cross, scheme=scheme)

    def to_dict(self) -> dict[str, object]:
        payload: dict[str, object] = {
            "lookahead": self.lookahead,
            "allow_cross_component_merge": self.allow_cross_component_merge,
        }

        if self.scheme != "set_bpe":
            payload["scheme"] = self.scheme
        return payload


@dataclass(frozen=True)
class PrimitiveTrainingConfig:
    """Serialized config-v1 metadata retained until artifacts are versioned again."""

    zero_primitive_components: tuple[str, ...]

    @classmethod
    def from_dict(
        cls,
        payload: object,
        *,
        physical: PhysicalConfig,
        path: str = "train.primitives",
    ) -> PrimitiveTrainingConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"zero_primitive_components"}, path)
        raw_names = _required(values, "zero_primitive_components", path)
        if not isinstance(raw_names, Sequence) or isinstance(raw_names, (str, bytes, bytearray)):
            raise ConfigError(f"{path}.zero_primitive_components must be a list")
        names = tuple(raw_names)
        if not all(isinstance(name, str) and name for name in names):
            raise ConfigError(
                f"{path}.zero_primitive_components must contain non-empty component names"
            )
        if len(set(names)) != len(names):
            raise ConfigError(f"{path}.zero_primitive_components contains a duplicate")
        unknown = sorted(set(names) - set(physical.components))
        if unknown:
            raise ConfigError(
                f"{path}.zero_primitive_components contains unknown component {unknown[0]!r}"
            )

        selected = set(names)
        for name, component in physical.components.items():
            if component.space == "additive":
                if component.primitive_order not in {0, 1, 2}:
                    raise ConfigError(
                        f"{path} cannot train additive primitive_order="
                        f"{component.primitive_order} for component {name!r}"
                    )
                domain_order = component.primitive_order
                if name in selected and component.primitive_order == 0:
                    raise ConfigError(
                        f"{path}.zero_primitive_components cannot contain absolute "
                        f"component {name!r}"
                    )
            else:
                domain_order = component.primitive_order
            if f"order_{domain_order}" not in component.constraints:
                raise ConfigError(
                    f"physical component {name!r} needs constraints.order_{domain_order} "
                    "to train its primitives"
                )
        return cls(zero_primitive_components=names)

    def to_dict(self) -> dict[str, object]:
        return {"zero_primitive_components": list(self.zero_primitive_components)}


@dataclass(frozen=True)
class BPETrainingConfig:
    """Spatial and final vocabulary budgets used for whole-trajectory BPE."""

    spatial_vocab_size: int
    vocab_size: int

    @classmethod
    def from_dict(cls, payload: object, *, path: str = "train.bpe") -> BPETrainingConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"spatial_vocab_size", "vocab_size"}, path)
        spatial_vocab_size = _integer(
            _required(values, "spatial_vocab_size", path),
            f"{path}.spatial_vocab_size",
            positive=True,
        )
        vocab_size = _integer(
            _required(values, "vocab_size", path), f"{path}.vocab_size", positive=True
        )
        if spatial_vocab_size > vocab_size:
            raise ConfigError(f"{path}.spatial_vocab_size must not exceed vocab_size")
        return cls(spatial_vocab_size=spatial_vocab_size, vocab_size=vocab_size)

    def to_dict(self) -> dict[str, object]:
        return {
            "spatial_vocab_size": self.spatial_vocab_size,
            "vocab_size": self.vocab_size,
        }


@dataclass(frozen=True)
class QualityControlConfig:
    """Pretraining-only policy for extreme first-order discontinuities."""

    hard_order_1_scale: float
    min_segment_length: int

    @classmethod
    def from_dict(
        cls, payload: object, *, path: str = "train.quality_control"
    ) -> QualityControlConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"hard_order_1_scale", "min_segment_length"}, path)
        scale = _number(
            _required(values, "hard_order_1_scale", path),
            f"{path}.hard_order_1_scale",
        )
        if scale <= 1.0:
            raise ConfigError(f"{path}.hard_order_1_scale must be greater than 1")
        return cls(
            hard_order_1_scale=scale,
            min_segment_length=_integer(
                _required(values, "min_segment_length", path),
                f"{path}.min_segment_length",
                positive=True,
            ),
        )

    def to_dict(self) -> dict[str, object]:
        return {
            "hard_order_1_scale": self.hard_order_1_scale,
            "min_segment_length": self.min_segment_length,
        }


@dataclass(frozen=True)
class SparseTrainingConfig:
    """ActionCodec2Sparse training-corpus construction."""

    chunk_horizon: int
    static_run_cap: int

    @classmethod
    def from_dict(cls, payload: object, *, path: str = "train.sparse") -> SparseTrainingConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"chunk_horizon", "static_run_cap"}, path)
        return cls(
            chunk_horizon=_integer(
                _required(values, "chunk_horizon", path), f"{path}.chunk_horizon", positive=True
            ),
            static_run_cap=_integer(
                _required(values, "static_run_cap", path), f"{path}.static_run_cap", positive=True
            ),
        )

    def to_dict(self) -> dict[str, object]:
        return {"chunk_horizon": self.chunk_horizon, "static_run_cap": self.static_run_cap}


@dataclass(frozen=True)
class TrainConfig:
    """Non-derivable fit choices and BPE budgets."""

    primitives: PrimitiveTrainingConfig
    bpe: BPETrainingConfig
    quality_control: QualityControlConfig | None = None
    sparse: SparseTrainingConfig | None = None

    @classmethod
    def from_dict(
        cls,
        payload: object,
        *,
        physical: PhysicalConfig,
        path: str = "train",
    ) -> TrainConfig:
        values = _mapping(payload, path)
        _reject_unknown(values, {"primitives", "bpe", "quality_control", "sparse"}, path)
        raw_quality_control = values.get("quality_control")
        raw_sparse = values.get("sparse")
        return cls(
            primitives=PrimitiveTrainingConfig.from_dict(
                _required(values, "primitives", path), physical=physical
            ),
            bpe=BPETrainingConfig.from_dict(_required(values, "bpe", path)),
            quality_control=(
                None
                if raw_quality_control is None
                else QualityControlConfig.from_dict(raw_quality_control)
            ),
            sparse=None if raw_sparse is None else SparseTrainingConfig.from_dict(raw_sparse),
        )

    def to_dict(self) -> dict[str, object]:
        payload = {"primitives": self.primitives.to_dict(), "bpe": self.bpe.to_dict()}
        if self.quality_control is not None:
            payload["quality_control"] = self.quality_control.to_dict()
        if self.sparse is not None:
            payload["sparse"] = self.sparse.to_dict()
        return payload


@dataclass(frozen=True)
class ActionCodec2Config:
    """Strict YAML configuration for an ActionCodec2 processor artifact."""

    physical: PhysicalConfig
    set_bpe: SetBPEConfig
    train: TrainConfig | None = None
    model_type: str = MODEL_TYPE
    config_version: int = CONFIG_VERSION

    @classmethod
    def from_dict(cls, payload: object) -> ActionCodec2Config:
        values = _mapping(payload, "config")
        _reject_unknown(
            values,
            {"model_type", "config_version", "physical", "set_bpe", "train"},
            "config",
        )
        model_type = _required(values, "model_type", "config")
        if model_type != MODEL_TYPE:
            raise ConfigError(f"config.model_type must be {MODEL_TYPE!r}")
        config_version = _integer(
            _required(values, "config_version", "config"), "config.config_version"
        )
        if config_version != CONFIG_VERSION:
            raise ConfigError(
                f"unsupported config_version {config_version}; expected {CONFIG_VERSION}"
            )
        physical = PhysicalConfig.from_dict(_required(values, "physical", "config"))
        raw_train = values.get("train")
        set_bpe = SetBPEConfig.from_dict(_required(values, "set_bpe", "config"))
        train = None if raw_train is None else TrainConfig.from_dict(raw_train, physical=physical)
        sparse = set_bpe.scheme == "sparse_set_bpe"
        if train is not None and (train.sparse is not None) != sparse:
            raise ConfigError(
                "train.sparse is required for, and only valid with, set_bpe.scheme=sparse_set_bpe"
            )
        return cls(
            physical=physical,
            set_bpe=set_bpe,
            train=train,
            model_type=str(model_type),
            config_version=config_version,
        )

    @classmethod
    def from_yaml(cls, path: str | Path) -> ActionCodec2Config:
        source = Path(path)
        try:
            payload = yaml.load(source.read_text(encoding="utf-8"), Loader=_UniqueKeyLoader)
        except OSError:
            raise
        except yaml.YAMLError as error:
            raise ConfigError(f"cannot parse {source}: {error}") from error
        return cls.from_dict(payload)

    @classmethod
    def from_pretrained(cls, directory: str | Path) -> ActionCodec2Config:
        """Load the strict numerical configuration from ``config.yaml``."""
        return cls.from_yaml(Path(directory) / CONFIG_NAME)

    def to_dict(self) -> dict[str, object]:
        payload: dict[str, object] = {
            "model_type": self.model_type,
            "config_version": self.config_version,
            "physical": self.physical.to_dict(),
            "set_bpe": self.set_bpe.to_dict(),
        }
        if self.train is not None:
            payload["train"] = self.train.to_dict()
        return payload

    def save_yaml(self, path: str | Path) -> None:
        destination = Path(path)
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_text(
            yaml.safe_dump(
                self.to_dict(),
                sort_keys=False,
                allow_unicode=True,
                default_flow_style=False,
            ),
            encoding="utf-8",
        )

    def save_pretrained(self, directory: str | Path) -> None:
        """Write ``config.yaml`` without embedding primitive or BPE arrays."""
        self.save_yaml(Path(directory) / CONFIG_NAME)