File size: 12,084 Bytes
7c09a0c ce32ca3 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c ce32ca3 7c09a0c 1bfc3f0 7c09a0c ce32ca3 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 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 | """Load a named ActionCodec2 variant or a standalone saved processor."""
from __future__ import annotations
import hashlib
import importlib
import importlib.util
import json
from pathlib import Path
import re
import sys
PROFILE_FILES = ("config.yaml", "primitives.npz", "bpe.json", "config.json", "processor_config.json")
SHARED_PATTERNS = (
"config.json", "processor_config.json", "processing_actioncodec2.py", "variants.json",
"runtime/**", "actioncodec2_kernels/**", "setup.py", "pyproject.toml", "LICENSE",
"README.md", "requirements.txt",
)
VARIANT_ALIASES = {
**{f"paper-{order}": f"paper-o{index}-5mm-4k-2bin-1p2ms"
for index, order in enumerate(("1st", "2nd"), 1)},
**{f"extended-{order}-{precision}-{size}k":
f"extended-o{index}-{precision}-{size}k-2bin-1p2ms"
for index, order in enumerate(("1st", "2nd"), 1)
for precision in ("1mm", "5mm", "10mm") for size in (1, 2, 4, 8, 16)
if precision != "1mm" or size != 1},
**{f"sparse-{dynamics}-{mode}": f"sparse-o1-5mm-4k-{mode}-{speed}"
for dynamics, speed in (("standard", "1p2ms"), ("high-dynamic", "2ms"))
for mode in ("2bin", "21bin")},
}
def resolve_variant_name(variant, variants):
"""Resolve either published spelling against an old or new catalog."""
if variant in variants:
return variant
canonical = VARIANT_ALIASES.get(variant, variant)
if canonical in variants:
return canonical
for legacy, name in VARIANT_ALIASES.items():
if name == canonical and legacy in variants:
return legacy
raise ValueError(f"select variant= from: {', '.join(sorted(variants))}")
def asset_inventory(directory: Path) -> dict[str, str]:
"""Hash the complete configuration and fitted assets for one variant."""
sparse = (directory / "sparse_router_config.yaml").is_file()
names = ["sparse_router_config.yaml" if sparse else "router_config.yaml", "config.json", "processor_config.json"]
profiles = [name for name in ("eef", "joint") if (directory / "profiles" / name).is_dir()]
if not profiles:
raise ValueError("variant has no fitted profiles")
names.extend(f"profiles/{profile}/{name}" for profile in profiles
for name in (tuple("sparse_bpe.json" if n == "bpe.json" else n for n in PROFILE_FILES) if sparse else PROFILE_FILES))
names.extend(f"profiles/{profile}/preprocessing.json" for profile in profiles
if (directory / "profiles" / profile / "preprocessing.json").is_file())
return {name: hashlib.sha256((directory / name).read_bytes()).hexdigest()
for name in sorted(names)}
def inventory_fingerprint(files: dict[str, str]) -> str:
"""Return the stable identity of a fitted asset inventory."""
return hashlib.sha256(json.dumps(files, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
def _check_inventory(directory: Path, record: dict) -> None:
files = asset_inventory(directory)
if files != record.get("files") or inventory_fingerprint(files) != record.get("fingerprint"):
raise ValueError("ActionCodec2 asset fingerprint mismatch; restore the complete selected variant")
class _VariantBinding:
@property
def variant(self):
return self._variant_origin.get("name")
@property
def variant_fingerprint(self):
return self._variant_origin.get("fingerprint")
def for_action_space(self, name):
codec = super().for_action_space(name)
codec._variant_origin = dict(self._variant_origin)
return codec
def fit(self, *args, **kwargs):
result = super().fit(*args, **kwargs)
self._variant_origin = {}
return result
def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):
directory = Path(save_directory)
include_training_report = kwargs.pop("include_training_report", False)
if (directory / "variants.json").is_file():
raise ValueError("save a selected variant to a separate directory, outside the catalog root")
if not push_to_hub and kwargs:
raise TypeError(f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}")
files = super().save_pretrained(directory, include_training_report=include_training_report)
inventory = asset_inventory(directory)
manifest = directory / "artifact_manifest.json"
manifest.write_text(json.dumps({
"format_version": 1,
"origin": self._variant_origin,
"files": inventory,
"fingerprint": inventory_fingerprint(inventory),
}, indent=2) + "\n", encoding="utf-8")
files.append(str(manifest))
if push_to_hub:
package = type(self).__module__.split(".frontend.", 1)[0]
importlib.import_module(f"{package}.backends.actioncodec2.hf").upload_pretrained_directory(directory, **kwargs)
return files
class ActionCodec2:
"""Hugging Face entry point for catalog and standalone artifacts."""
@classmethod
def register_for_auto_class(cls, auto_class="AutoProcessor"):
return cls
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *, variant=None,
variant_fingerprint=None, action_space=None, **kwargs):
source = Path(pretrained_model_name_or_path).expanduser()
kwargs.pop("_from_auto", None)
kwargs.pop("trust_remote_code", None)
subfolder = kwargs.pop("subfolder", "")
if subfolder:
raise ValueError("select an ActionCodec2 catalog variant with variant=, not 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 variant is not None and (not isinstance(variant, str) or
re.fullmatch(r"[a-z0-9]+(?:-[a-z0-9]+)*", variant) is None):
raise ValueError("variant must be a published variant name")
if variant_fingerprint is not None and (
not isinstance(variant_fingerprint, str) or
re.fullmatch(r"[0-9a-f]{64}", variant_fingerprint) is None
):
raise ValueError("variant_fingerprint must be a SHA-256 fingerprint")
remote = not source.is_dir()
if remote and str(pretrained_model_name_or_path).startswith(("/", "./", "../", "~")):
raise FileNotFoundError(f"ActionCodec2 artifact does not exist: {source}")
revision = None
if remote:
from huggingface_hub import hf_hub_download, snapshot_download
code_revision = Path(__file__).parent.name
if re.fullmatch(r"[0-9a-f]{40}", code_revision):
requested = kwargs.get("revision")
if requested and re.fullmatch(r"[0-9a-f]{40}", requested) and requested != code_revision:
raise ValueError("ActionCodec2 code and assets must use the same revision")
kwargs["revision"] = code_revision
metadata_file = Path(hf_hub_download(
str(pretrained_model_name_or_path), "processor_config.json", **kwargs
))
root = metadata_file.parent
revision = root.name
kwargs["revision"] = revision
else:
root = source
metadata_file = root / "processor_config.json"
metadata = json.loads(metadata_file.read_text(encoding="utf-8"))
if not isinstance(metadata, dict):
raise ValueError("invalid ActionCodec2 processor configuration")
record = None
origin = {}
if metadata.get("variant_catalog") == "variants.json":
catalog_path = (Path(hf_hub_download(
str(pretrained_model_name_or_path), "variants.json", **kwargs
)) if remote else root / "variants.json")
catalog = json.loads(catalog_path.read_text(encoding="utf-8"))
if (not isinstance(catalog, dict) or catalog.get("format_version") != 1 or
not isinstance(catalog.get("variants"), dict)):
raise ValueError("unsupported ActionCodec2 variant catalog")
variants = catalog["variants"]
variant = resolve_variant_name(variant, variants)
record = variants[variant]
if not isinstance(record, dict):
raise ValueError("invalid ActionCodec2 variant record")
directory = root / "variants" / variant
if remote:
root = Path(snapshot_download(
str(pretrained_model_name_or_path),
allow_patterns=[*SHARED_PATTERNS, f"variants/{variant}/**"], **kwargs,
))
directory = root / "variants" / variant
_check_inventory(directory, record)
origin = {"name": variant, "fingerprint": record["fingerprint"]}
if remote:
origin.update(repo_id=str(pretrained_model_name_or_path), revision=revision)
else:
if remote:
root = Path(snapshot_download(
str(pretrained_model_name_or_path),
allow_patterns=[*SHARED_PATTERNS, "router_config.yaml", "profiles/**",
"sparse_router_config.yaml", "artifact_manifest.json"], **kwargs,
))
directory = root
manifest = root / "artifact_manifest.json"
if manifest.is_file():
record = json.loads(manifest.read_text(encoding="utf-8"))
if (not isinstance(record, dict) or record.get("format_version") != 1 or
not isinstance(record.get("origin"), dict)):
raise ValueError("unsupported ActionCodec2 saved artifact manifest")
_check_inventory(directory, record)
origin = record["origin"]
if variant is not None and VARIANT_ALIASES.get(variant, variant) != VARIANT_ALIASES.get(origin.get("name"), origin.get("name")):
raise ValueError("saved processor contains a different variant; load the catalog to switch")
if variant_fingerprint is not None and variant_fingerprint != origin.get("fingerprint"):
raise ValueError("variant fingerprint differs from the tokenizer bound to this policy")
family = "actioncodec2_sparse" if (directory / "sparse_router_config.yaml").is_file() else "actioncodec2"
if __package__ == "actioncodec2.core":
package_name = __package__
else:
runtime = root / "runtime"
package_file = runtime / "__init__.py"
if not package_file.is_file():
raise FileNotFoundError(f"ActionCodec2 runtime is missing from {root}")
identity = hashlib.sha256(str(runtime.resolve()).encode()).hexdigest()[:16]
package_name = f"_actioncodec2_artifact_{identity}"
if package_name not in sys.modules:
spec = importlib.util.spec_from_file_location(
package_name, package_file, submodule_search_locations=[str(runtime)]
)
package = importlib.util.module_from_spec(spec)
sys.modules[package_name] = package
spec.loader.exec_module(package)
processor = importlib.import_module(f"{package_name}.frontend.{family}.processor")
base = processor.ActionCodec2Sparse if family == "actioncodec2_sparse" else processor.ActionCodec2
if not hasattr(processor, "_catalog_processor_class"):
processor._catalog_processor_class = type(
base.__name__, (_VariantBinding, base), {"__module__": processor.__name__}
)
codec = processor._catalog_processor_class.from_pretrained(directory, action_space=action_space)
codec._variant_origin = origin
codec._source_directory = root
return codec
|