File size: 7,274 Bytes
1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | '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
|