| """Package-owned configuration templates for fresh ActionCodec2 instances.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Mapping |
| from functools import lru_cache |
| from pathlib import Path |
| from types import MappingProxyType |
|
|
| from .bpe import BPESerializer, SetBPESerializer |
| from .configuration import ActionCodec2Config |
| from .presets import PresetConfig |
| from .processor import ActionCodec2Processor |
| from .routing import ActionCodec2Router |
| from .training import build_primitives |
|
|
|
|
| class CodecBuilder: |
| 'Build unfitted physical profiles from packaged defaults or explicit configs.' |
|
|
| def __init__(self, *, primitive_order=1, bpe_scheme=None, profile_configs=None): |
| if type(primitive_order) is not int or primitive_order not in (1, 2): |
| raise ValueError("primitive_order must be 1 or 2") |
| if bpe_scheme not in (None, "set_bpe", "unified_set_bpe", "bpe"): |
| raise ValueError("bpe_scheme must be 'set_bpe', 'unified_set_bpe' or 'bpe'") |
| if profile_configs is not None and not isinstance(profile_configs, Mapping): |
| raise TypeError("profile_configs must map 'joint'/'eef' to physical configurations") |
| self.profile_configs = dict(profile_configs or {}) |
| if set(self.profile_configs) - {"joint", "eef"}: |
| raise ValueError("profile_configs keys must be 'joint' or 'eef'") |
| self.primitive_order = primitive_order |
| self.bpe_scheme = bpe_scheme |
| self.assets = Path(__file__).with_name("assets") |
|
|
| def _config(self, name: str) -> ActionCodec2Config: |
| suffix = "" if self.primitive_order == 1 else "_second_order" |
| raw = self.profile_configs.get(name, self.assets / "profiles" / f"{name}{suffix}.yaml") |
| if isinstance(raw, ActionCodec2Config): |
| config = raw |
| elif isinstance(raw, Mapping): |
| config = ActionCodec2Config.from_dict(raw) |
| else: |
| config = ActionCodec2Config.from_yaml(raw) |
| if config.train is None: |
| raise ValueError(f"{name} profile config requires a train section") |
| if self.bpe_scheme is not None: |
| payload = config.to_dict() |
| payload["set_bpe"]["scheme"] = self.bpe_scheme |
| if self.bpe_scheme == "bpe": |
| payload["train"]["bpe"]["spatial_vocab_size"] = payload["train"]["bpe"][ |
| "vocab_size" |
| ] |
| config = ActionCodec2Config.from_dict(payload) |
| return config |
|
|
| def build(self) -> ActionCodec2Router: |
| """Create primitive-only profiles; the public facade tracks fitted state.""" |
| processors = {} |
| budgets = {} |
| for name in ("joint", "eef"): |
| config = self._config(name) |
| primitives = build_primitives(config) |
| if config.set_bpe.scheme == "bpe": |
| serializer = BPESerializer( |
| bins=primitives.bins, lookahead=config.set_bpe.lookahead, merges=() |
| ) |
| elif config.set_bpe.scheme == "unified_set_bpe": |
| from .unified_set_bpe import UnifiedSetBPESerializer |
|
|
| serializer = UnifiedSetBPESerializer( |
| bins=primitives.bins, lookahead=config.set_bpe.lookahead, merges=() |
| ) |
| else: |
| serializer = SetBPESerializer( |
| bins=primitives.bins, |
| lookahead=config.set_bpe.lookahead, |
| spatial_merges=(), |
| temporal_merges=(), |
| ) |
| processors[name] = ActionCodec2Processor(config, primitives, serializer) |
| budgets[name] = config.train.bpe.vocab_size |
| presets = list(builtin_action_spaces().values()) |
| return ActionCodec2Router.from_processors( |
| processors, token_budgets=budgets, presets=presets |
| ) |
|
|
|
|
| @lru_cache(maxsize=1) |
| def builtin_action_spaces(): |
| 'Return a read-only name-to-preset mapping shipped in the installed package.' |
| presets = [ |
| PresetConfig.from_yaml(path) |
| for path in sorted(Path(__file__).with_name("assets").joinpath("presets").glob("*.yaml")) |
| ] |
| return MappingProxyType({preset.name: preset for preset in presets}) |
|
|
|
|
| class ProfileBundleBuilder: |
| 'Assemble two existing physical artifacts using their recorded vocabulary budgets.' |
|
|
| def __init__(self, joint, eef, *, token_budgets=None): |
| self.paths = {"joint": joint, "eef": eef} |
| self.budgets = dict(token_budgets or {}) |
| if self.budgets.keys() - self.paths.keys(): |
| raise ValueError("token_budgets accepts only 'joint' and 'eef'") |
|
|
| def build(self): |
| processors = { |
| name: ActionCodec2Processor.from_pretrained(path) for name, path in self.paths.items() |
| } |
| budgets = { |
| name: self.budgets.get( |
| name, |
| processor.config.train.bpe.vocab_size |
| if processor.config.train is not None |
| else processor.vocab_size, |
| ) |
| for name, processor in processors.items() |
| } |
| return ActionCodec2Router.from_processors( |
| processors, token_budgets=budgets, presets=list(builtin_action_spaces().values()) |
| ) |
|
|
