| """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) |
|
|
