| 'Hugging Face AutoProcessor integration.' |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import TYPE_CHECKING, Any |
|
|
| from transformers import AutoConfig, AutoProcessor, PretrainedConfig |
|
|
| from .configuration import CONFIG_NAME, MODEL_TYPE |
|
|
| if TYPE_CHECKING: |
| from .processor import ActionCodec2Processor |
| from .routing import ActionCodec2Router |
|
|
| HF_CONFIG_NAME = "config.json" |
| PROCESSOR_CONFIG_NAME = "processor_config.json" |
| PROCESSOR_CLASS_NAME = "ActionCodec2Processor" |
| ROUTER_MODEL_TYPE = "actioncodec2_router" |
| ROUTER_PROCESSOR_CLASS_NAME = "ActionCodec2Router" |
| ROUTER_CONFIG_NAME = "router_config.yaml" |
| GENERATED_CARD_MARKER = ( |
| "<!-- generated by actioncodec2; delete this line to keep your own edits -->" |
| ) |
|
|
|
|
| ROUTER_MANIFEST_PATTERNS = [ |
| ROUTER_CONFIG_NAME, |
| HF_CONFIG_NAME, |
| PROCESSOR_CONFIG_NAME, |
| f"profiles/*/{CONFIG_NAME}", |
| "manifest.json", |
| f"joint/**/{CONFIG_NAME}", |
| f"eef/**/{CONFIG_NAME}", |
| ] |
| ROUTER_ALLOW_PATTERNS = [ |
| *ROUTER_MANIFEST_PATTERNS, |
| "profiles/**", |
| "joint/**", |
| "eef/**", |
| "fit_report.json", |
| "README.md", |
| ] |
|
|
|
|
| class ActionCodec2FrontendHFConfig(PretrainedConfig): |
| """AutoConfig discovery key for the action-space-bound public processor.""" |
|
|
| model_type = "actioncodec2" |
|
|
| def __init__(self, **kwargs: Any) -> None: |
| kwargs.setdefault("processor_class", "ActionCodec2") |
| super().__init__(**kwargs) |
|
|
|
|
| class ActionCodec2HFConfig(PretrainedConfig): |
| """Minimal AutoConfig key that dispatches to :class:`ActionCodec2Processor`.""" |
|
|
| model_type = MODEL_TYPE |
|
|
| def __init__(self, **kwargs: Any) -> None: |
| kwargs.setdefault("processor_class", PROCESSOR_CLASS_NAME) |
| super().__init__(**kwargs) |
|
|
|
|
| class ActionCodec2RouterHFConfig(PretrainedConfig): |
| """Minimal AutoConfig key that dispatches to :class:`ActionCodec2Router`.""" |
|
|
| model_type = ROUTER_MODEL_TYPE |
|
|
| def __init__(self, **kwargs: Any) -> None: |
| kwargs.setdefault("processor_class", ROUTER_PROCESSOR_CLASS_NAME) |
| super().__init__(**kwargs) |
|
|
|
|
| def save_huggingface_metadata( |
| directory: str | Path, |
| *, |
| model_type: str = MODEL_TYPE, |
| processor_class: str = PROCESSOR_CLASS_NAME, |
| **processor_config: Any, |
| ) -> tuple[Path, Path]: |
| """Write standard local/Hub discovery files for ``AutoProcessor``.""" |
| root = Path(directory) |
| root.mkdir(parents=True, exist_ok=True) |
| config_path = root / HF_CONFIG_NAME |
| processor_path = root / PROCESSOR_CONFIG_NAME |
| config_path.write_text( |
| json.dumps( |
| { |
| "model_type": model_type, |
| "processor_class": processor_class, |
| }, |
| indent=2, |
| ) |
| + "\n", |
| encoding="utf-8", |
| ) |
| processor_path.write_text( |
| json.dumps({"processor_class": processor_class, **processor_config}, indent=2) + "\n", |
| encoding="utf-8", |
| ) |
| return config_path, processor_path |
|
|
|
|
| def resolve_pretrained_directory( |
| pretrained_model_name_or_path: str | Path, |
| *, |
| allow_patterns: list[str] | None = None, |
| **kwargs: Any, |
| ) -> Path: |
| """Resolve a local directory or download one immutable Hub snapshot.""" |
| source = Path(pretrained_model_name_or_path) |
| kwargs.pop("_from_auto", None) |
| kwargs.pop("trust_remote_code", None) |
| subfolder = str(kwargs.pop("subfolder", "")) |
| allowed = { |
| "cache_dir", |
| "force_download", |
| "local_files_only", |
| "token", |
| "revision", |
| "repo_type", |
| } |
| unknown = sorted(set(kwargs) - allowed) |
| if unknown: |
| raise TypeError(f"unsupported from_pretrained keyword(s): {', '.join(unknown)}") |
| if source.is_dir(): |
| return source / subfolder |
|
|
| try: |
| from huggingface_hub import snapshot_download |
| except ImportError as error: |
| raise ImportError( |
| "loading a Hub artifact requires huggingface_hub (installed with Transformers)" |
| ) from error |
| patterns = allow_patterns or [ |
| "config.yaml", |
| "primitives.npz", |
| "bpe.json", |
| HF_CONFIG_NAME, |
| PROCESSOR_CONFIG_NAME, |
| ] |
| if subfolder: |
| prefix = subfolder.rstrip("/") |
| patterns = [f"{prefix}/{pattern}" for pattern in patterns] |
| snapshot = snapshot_download( |
| repo_id=str(pretrained_model_name_or_path), |
| allow_patterns=patterns, |
| **kwargs, |
| ) |
| return Path(snapshot) / subfolder |
|
|
|
|
| def upload_pretrained_directory(directory: Path, **kwargs: Any) -> None: |
| """Upload a complete, already-saved processor artifact to the Hub.""" |
| from huggingface_hub import HfApi |
|
|
| repo_id = kwargs.pop("repo_id", directory.name) |
| commit_message = kwargs.pop("commit_message", "Upload ActionCodec2 processor") |
| api = HfApi(token=kwargs.pop("token", None)) |
| api.create_repo(repo_id=repo_id, exist_ok=True, **kwargs) |
| api.upload_folder( |
| repo_id=repo_id, |
| folder_path=str(directory), |
| commit_message=commit_message, |
| ) |
|
|
|
|
| def register_actioncodec2_auto_processors( |
| processor_class: type[ActionCodec2Processor], |
| router_class: type[ActionCodec2Router], |
| ) -> bool: |
| 'Register both installed ActionCodec2 artifact types with Transformers.' |
| AutoConfig.register(MODEL_TYPE, ActionCodec2HFConfig, exist_ok=True) |
| AutoProcessor.register(ActionCodec2HFConfig, processor_class, exist_ok=True) |
| AutoConfig.register(ROUTER_MODEL_TYPE, ActionCodec2RouterHFConfig, exist_ok=True) |
| AutoProcessor.register(ActionCodec2RouterHFConfig, router_class, exist_ok=True) |
| return True |
|
|
|
|
| def save_model_card(directory: str | Path, description: str) -> Path: |
| 'Write the generated Hub card, refreshing it whenever the artifact is re-saved.' |
| path = Path(directory) / "README.md" |
| if not path.exists() or GENERATED_CARD_MARKER in path.read_text(encoding="utf-8"): |
| path.parent.mkdir(parents=True, exist_ok=True) |
| guide = ( |
| "---\ntags:\n- robotics\n- action-tokenizer\n---\n\n" |
| f"{GENERATED_CARD_MARKER}\n\n" |
| "# ActionCodec2\n\n" |
| "This artifact contains physical vocabularies, BPE rules and action-space layouts.\n" |
| "Install actioncodec2 and load the saved processor locally:\n\n" |
| "```python\nfrom actioncodec2 import ActionCodec2\n" |
| "from transformers import AutoProcessor\n\n" |
| "print(ActionCodec2.describe_pretrained('./my_codec'))\n" |
| "codec = AutoProcessor.from_pretrained(\n" |
| " './my_codec', action_space='single_eef_delta',\n)\n" |
| "print(codec)\n```\n\n" |
| "Choose action_space from your data's physical meaning, units and column order;\n" |
| "the example layout is not an automatic selection. For a bound artifact, omit\n" |
| "action_space to restore its saved binding. Pass your actual sampling rate as\n" |
| "fps to encode/decode. Absolute layouts require current_state; second-order\n" |
| "profiles require previous_first_order. See the artifact's layout guide below.\n" |
| ) |
| path.write_text( |
| guide + "\n\n## 当前 artifact 的实际布局与要求\n\n```text\n" + description + "\n```\n", |
| encoding="utf-8", |
| ) |
| return path |
|
|
