File size: 16,139 Bytes

"""Hugging Face processor over the shared ActionCodec2 frontend implementation."""

from __future__ import annotations

import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import numpy as np
import torch
from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin

from ...backends.actioncodec2 import ActionCodec2Router, PresetConfig
from ...backends.actioncodec2.defaults import (
    CodecBuilder,
    ProfileBundleBuilder,
    builtin_action_spaces,
)
from ...backends.actioncodec2.hf import (
    PROCESSOR_CONFIG_NAME,
    ROUTER_ALLOW_PATTERNS,
    ROUTER_MANIFEST_PATTERNS,
    ActionCodec2FrontendHFConfig,
    resolve_pretrained_directory,
    save_huggingface_metadata,
    save_model_card,
    upload_pretrained_directory,
)
from ...backends.actioncodec2.preprocessing import PreprocessingSpec
from ...backends.actioncodec2.release import MANIFEST_NAME
from ...exporting import save_runtime
from .loading import describe_release, load_release_router
from .tokenizer import ActionCodec2Tokenizer
from ..training import TrainingMixin
from ...backends.actioncodec2.training import ProfileTrainer


def _saved_metadata(directory: Path) -> dict:
    metadata = directory / PROCESSOR_CONFIG_NAME
    return json.loads(metadata.read_text(encoding="utf-8")) if metadata.is_file() else {}


def _fitted_status(profiles) -> str:
    status = ", ".join(sorted(profiles)) or "none; call codec.fit(episodes)"
    return f"Fitted profiles: {status}\n"


class ActionCodec2(TrainingMixin, ActionCodec2Tokenizer, ProcessorMixin):
    'Action-space-bound processor with lossless token serialization.'

    trainer_class = ProfileTrainer
    model_type = ActionCodec2FrontendHFConfig.model_type
    model_input_names = ["input_ids"]

    attributes = []

    def __init__(
        self,
        backend: ActionCodec2Router | None = None,
        *,
        action_space: str,
        primitive_order: int = 1,
        bpe_scheme: str | None = None,
        profile_configs: Mapping[str, object] | None = None,
    ) -> None:
        fresh = backend is None
        if fresh:
            backend = CodecBuilder(
                primitive_order=primitive_order,
                bpe_scheme=bpe_scheme,
                profile_configs=profile_configs,
            ).build()
        elif primitive_order != 1 or bpe_scheme is not None or profile_configs is not None:
            raise ValueError("profile construction options cannot accompany an existing backend")
        super().__init__(backend, action_space=action_space)
        ProcessorMixin.__init__(self)
        self._fitted_profiles = frozenset() if fresh else frozenset(backend.profiles)
        self.training_report = None
        self.preprocessing: dict[str, PreprocessingSpec] = {}
        self.release_metadata: dict[str, Any] | None = None

    def __dir__(self):
        inapplicable = {
            "audio_tokenizer",
            "chat_template",
            "feature_extractor_class",
            "apply_chat_template",
            "post_process_image_text_to_text",
            "pad_action_id",
            "bos_blk_id",
            "parts_meta",
            "BAR_SENTINEL_TOKENS",
        }
        return [name for name in super().__dir__() if name not in inapplicable]

    def __repr__(self) -> str:
        return (
            f"ActionCodec2(action_space={self.action_space.name!r}, "
            f"D={self.action_dim}, vocab={self.action_tokenizer.vocab_size})"
        )

    @staticmethod
    def builtin_action_spaces():
        """Return the immutable named layout definitions shipped in the package."""
        return builtin_action_spaces()

    @classmethod
    def from_profiles(cls, joint, eef, *, action_space: str, token_budgets=None):
        'Bind an artifact assembled from existing joint/EEF profile directories.'
        router = ProfileBundleBuilder(joint, eef, token_budgets=token_budgets).build()
        return cls(router, action_space=action_space)

    def logits_processor(
        self,
        horizon: int,
        *,
        fps: float,
        prompt_length: int,
        eos_token_id: int,
        token_offset: int = 0,
        pad_token_id=None,
    ):
        'Create a Transformers LogitsProcessor for a fixed action horizon.'
        from .generation import ActionCodec2LogitsProcessor

        return ActionCodec2LogitsProcessor(
            self,
            horizon,
            fps=fps,
            prompt_length=prompt_length,
            eos_token_id=eos_token_id,
            token_offset=token_offset,
            pad_token_id=pad_token_id,
        )

    def batch_decode(
        self, token_rows, *, fps=None, current_state=None, previous_first_order=None, errors="raise"
    ):
        'Decode independent token rows, including different valid horizons.'
        self._require_fitted()
        if errors not in ("raise", "return"):
            raise ValueError("errors must be 'raise' or 'return'")
        rows = list(token_rows)
        current = (
            None
            if current_state is None
            else np.asarray(
                current_state.detach().cpu()
                if isinstance(current_state, torch.Tensor)
                else current_state
            )
        )
        if current is not None and current.shape != (len(rows), self.action_dim):
            raise ValueError(f"current_state must have shape {(len(rows), self.action_dim)}")
        previous = self._numpy_boundary(previous_first_order)
        if previous is not None and any(
            np.ndim(v) != 2 or len(v) != len(rows) for v in previous.values()
        ):
            raise ValueError(
                "previous_first_order arrays must have batch dimension equal to token rows"
            )
        result = []
        for index, row in enumerate(rows):
            try:
                decoded = self.decode(
                    [row],
                    fps=fps,
                    current_state=None if current is None else current[index : index + 1],
                    previous_first_order=None
                    if previous is None
                    else {k: v[index : index + 1] for k, v in previous.items()},
                )
            except ValueError as error:
                error = ValueError(f"token row {index}: {error}")
                if errors == "raise":
                    raise error from None
                decoded = error
            result.append(decoded)
        return result

    def apply_chat_template(self, *args, **kwargs):
        raise NotImplementedError(
            "ActionCodec2 accepts action tensors; use codec.encode(actions, fps=...) or codec(actions, fps=...)"
        )

    def post_process_image_text_to_text(self, *args, **kwargs):
        raise NotImplementedError(
            "ActionCodec2 returns actions; use codec.decode(token_rows, fps=...)"
        )

    @property
    def fitted_profiles(self) -> frozenset[str]:
        """Names of profiles fitted on data; persisted by save_pretrained()."""
        return self._fitted_profiles

    @property
    def is_fitted(self) -> bool:
        """Whether every profile used by the selected action space is fitted."""
        return all(c.space in self.fitted_profiles for c in self.action_space.components)

    def _require_fitted(self):
        missing = {c.space for c in self.action_space.components} - self.fitted_profiles
        if missing:
            raise RuntimeError(
                f"Unfitted profiles {sorted(missing)}; call codec.fit(episodes) first"
            )

    def _describe_action_spaces(self, action_dim: int | None = None) -> str:
        return _fitted_status(self.fitted_profiles) + super()._describe_action_spaces(action_dim)

    def for_action_space(self, name: str) -> ActionCodec2:
        """Share current weights with a new binding; subsequent fit replaces only that instance."""
        codec = super().for_action_space(name)
        codec._fitted_profiles = self.fitted_profiles
        codec.training_report = self.training_report
        codec.preprocessing = dict(self.preprocessing)
        codec.release_metadata = self.release_metadata
        return codec


    def _encode_action_indices(self, action, encode_kwargs=None):
        self._require_fitted()
        return super()._encode_action_indices(action, encode_kwargs)

    def _decode_action_indices(self, action_indices, **kwargs):
        self._require_fitted()
        return super()._decode_action_indices(action_indices, **kwargs)

    def decode_prefix(self, action_token_ids, **kwargs):
        self._require_fitted()
        return super().decode_prefix(action_token_ids, **kwargs)

    def grammar(self, horizon: int, *, fps: float | None = None):
        self._require_fitted()
        return super().grammar(horizon, fps=fps)

    def final_first_order(self, action_token_ids, *, previous_first_order, executed_steps=None):
        self._require_fitted()
        return super().final_first_order(
            action_token_ids,
            previous_first_order=previous_first_order,
            executed_steps=executed_steps,
        )

    @classmethod
    def describe_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        *,
        action_dim: int | None = None,
        gripper_mode: str = "2bin",
        primitive_order: int = 1,
        **kwargs: Any,
    ) -> str:
        'Inspect a local/Hub artifact before choosing an action space.'
        directory = resolve_pretrained_directory(
            pretrained_model_name_or_path, allow_patterns=ROUTER_MANIFEST_PATTERNS, **kwargs
        )
        if (directory / MANIFEST_NAME).is_file():
            return describe_release(
                directory,
                action_dim=action_dim,
                gripper_mode=gripper_mode,
                primitive_order=primitive_order,
            )
        specs, _ = ActionCodec2Router.read_manifest(directory)
        metadata = _saved_metadata(directory)
        fitted = metadata.get("fitted_profiles", [spec.name for spec in specs])
        return _fitted_status(fitted) + ActionCodec2Router.describe_pretrained(
            directory, selected=metadata.get("action_space"), action_dim=action_dim
        )

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        *,
        action_space: str | Path | PresetConfig | Mapping[str, object] | None = None,
        gripper_mode: str | None = None,
        primitive_order: int | None = None,
        verify_gates: bool = False,
        **kwargs: Any,
    ) -> ActionCodec2:
        'Load a local/Hub artifact and restore or override its action space.'
        if cls is ActionCodec2 and ("variant" in kwargs or (Path(pretrained_model_name_or_path) / "variants.json").is_file()
                                   or (Path(pretrained_model_name_or_path) / "artifact_manifest.json").is_file()):
            from ...catalog import ActionCodec2 as Catalog
            return Catalog.from_pretrained(pretrained_model_name_or_path, action_space=action_space, **kwargs)
        directory = resolve_pretrained_directory(
            pretrained_model_name_or_path, allow_patterns=ROUTER_ALLOW_PATTERNS, **kwargs
        )
        if action_space is None:
            action_space = _saved_metadata(directory).get("action_space")
        if action_space is None:
            raise ValueError(
                "action_space is required when loading an unbound router artifact. "
                "First inspect its layouts with print(ActionCodec2.describe_pretrained(path)), "
                "then pass action_space='NAME'; shape alone cannot identify physical meaning."
            )
        if (directory / MANIFEST_NAME).is_file():
            router, preprocessing, metadata = load_release_router(
                directory,
                gripper_mode="2bin" if gripper_mode is None else gripper_mode,
                primitive_order=1 if primitive_order is None else primitive_order,
                verify_gates=verify_gates,
            )
            codec = cls._from_router(router, action_space=action_space)
            codec.preprocessing, codec.release_metadata = preprocessing, metadata
            return codec
        if gripper_mode is not None or primitive_order is not None or verify_gates:
            raise ValueError(
                "gripper_mode, primitive_order and verify_gates select release variants only"
            )
        codec = super().from_pretrained(directory, action_space=action_space)
        metadata = _saved_metadata(directory)
        fitted = metadata.get("fitted_profiles", list(codec.action_tokenizer.profiles))
        if not isinstance(fitted, list) or any(
            not isinstance(name, str) or name not in codec.action_tokenizer.profiles
            for name in fitted
        ):
            raise ValueError("invalid fitted_profiles in processor_config.json")
        codec._fitted_profiles = frozenset(fitted)
        policies = metadata.get("preprocessing", {})
        if (
            not isinstance(policies, dict)
            or policies.keys() - codec.action_tokenizer.profiles.keys()
        ):
            raise ValueError("invalid preprocessing profiles in processor_config.json")
        codec.preprocessing = {
            name: PreprocessingSpec.from_dict(spec) for name, spec in policies.items()
        }
        codec.release_metadata = metadata.get("release")
        report_path = directory / "fit_report.json"
        if report_path.is_file():
            codec.training_report = json.loads(report_path.read_text(encoding="utf-8"))
        return codec

    def to_dict(self, legacy_serialization: bool = True) -> dict[str, object]:
        'Return lightweight discovery metadata.'
        del legacy_serialization
        payload = {
            "processor_class": type(self).__name__,
            "model_type": self.model_type,
            "action_space": self.action_space.name,
            "fitted_profiles": sorted(self.fitted_profiles),
        }
        if self.preprocessing:
            payload["preprocessing"] = {
                name: spec.to_dict() for name, spec in self.preprocessing.items()
            }
        return payload

    def save_pretrained(
        self, save_directory: str | Path, push_to_hub: bool = False, *, include_training_report=False, **kwargs: Any
    ) -> list[str]:
        """Save the complete router and the selected action-space binding."""
        if not push_to_hub and kwargs:
            raise TypeError(f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}")
        directory = Path(save_directory)
        files = self.action_tokenizer.save_pretrained(directory)

        save_model_card(directory, str(self))
        save_huggingface_metadata(
            directory,
            **self.to_dict(),
        )
        report_path = directory / "fit_report.json"
        if include_training_report and self.training_report is not None:
            report_path.write_text(
                json.dumps(self.training_report, indent=2) + "\n", encoding="utf-8"
            )
            files.append(str(report_path))
        elif report_path.exists():
            report_path.unlink()
        files.extend(save_runtime(directory))
        if push_to_hub:
            upload_pretrained_directory(directory, **kwargs)
        return files

    def __call__(
        self,
        action: torch.Tensor | np.ndarray,
        *,
        fps: float | None = None,
        current_state: torch.Tensor | np.ndarray | None = None,
        previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None,
        encode_kwargs: dict | None = None,
        return_tensors: str | None = None,
    ) -> BatchFeature:
        
        rows = self.encode(
            action,
            fps=fps,
            current_state=current_state,
            previous_first_order=previous_first_order,
            encode_kwargs=encode_kwargs,
        )
        return BatchFeature({"input_ids": rows}, tensor_type=return_tensors)