"""Strict, pickle-free loader for an AF3-NA+ Hugging Face directory. The loader reads only the model configuration, tokenizer contract, and ``model.safetensors`` weights needed for inference. Release audit metadata is kept outside the runtime path. """ from __future__ import annotations from dataclasses import dataclass import json import pathlib from types import MappingProxyType from typing import Any, Mapping class HubLoadError(ValueError): """Raised before a release is allowed to allocate or run a model.""" EXPECTED_VOCABULARY = ( "A", "C", "G", "U", "T", "N", "MASK", "BOS", "EOS", "PAD", "CHAIN_BOUNDARY", ) @dataclass(frozen=True) class ReleaseSpec: root: pathlib.Path config: Mapping[str, Any] tokenizer: Mapping[str, Any] model_config: Any precision: str @dataclass(frozen=True) class LoadedRelease: model: Any spec: ReleaseSpec device: Any dtype: Any def _reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: result: dict[str, Any] = {} for key, value in pairs: if key in result: raise HubLoadError(f"duplicate JSON key {key!r}") result[key] = value return result def _load_json_object(path: pathlib.Path) -> dict[str, Any]: try: value = json.loads( path.read_text(encoding="utf-8"), object_pairs_hook=_reject_duplicate_keys ) except OSError as exc: raise HubLoadError(f"cannot read {path.name}: {exc}") from exc except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise HubLoadError(f"invalid {path.name}: {exc}") from exc if not isinstance(value, dict): raise HubLoadError(f"{path.name} root must be an object") return value def _validate_tokenizer(tokenizer: Mapping[str, Any], *, max_chains: int) -> None: required = { "schema_version", "max_chains", "vocabulary", "one_token_per_nucleotide", "explicit_chain_boundary", "reject_t_in_rna", "reject_u_in_dna", "silent_ut_conversion", } if set(tokenizer) != required: raise HubLoadError( "tokenizer keys differ from the AF3-NA+ schema: " f"missing={sorted(required - set(tokenizer))} " f"extra={sorted(set(tokenizer) - required)}" ) if tokenizer["schema_version"] != 1: raise HubLoadError("unsupported tokenizer schema_version") if tokenizer["max_chains"] != max_chains: raise HubLoadError("tokenizer max_chains differs from model config") if tuple(tokenizer["vocabulary"]) != EXPECTED_VOCABULARY: raise HubLoadError("tokenizer vocabulary/order differs from the model contract") required_true = ( "one_token_per_nucleotide", "explicit_chain_boundary", "reject_t_in_rna", "reject_u_in_dna", ) if any(tokenizer[name] is not True for name in required_true): raise HubLoadError("required tokenizer safety rule is disabled") if tokenizer["silent_ut_conversion"] is not False: raise HubLoadError("silent U/T conversion must remain disabled") def load_release_spec(model_dir: str | pathlib.Path) -> ReleaseSpec: """Parse and validate release metadata without loading model weights.""" root = pathlib.Path(model_dir).expanduser().resolve(strict=True) if not root.is_dir(): raise HubLoadError("model_dir must be a directory") config = _load_json_object(root / "config.json") tokenizer = _load_json_object(root / "tokenizer.json") required = { "schema_version", "model_type", "precision", "model_config", "checkpoint_file", "quantization", } missing = required - set(config) if missing: raise HubLoadError(f"config.json is missing {sorted(missing)}") if config["schema_version"] != 1 or config["model_type"] != "evotemplate_na": raise HubLoadError("unsupported AF3-NA+ config schema/model_type") if config["precision"] != "bfloat16" or config["quantization"] is not None: raise HubLoadError("this runtime requires the BF16, non-quantized release") if config["checkpoint_file"] != "model.safetensors": raise HubLoadError("checkpoint_file must be model.safetensors") from evotemplate_na.model.config import EvoTemplateConfig try: model_config = EvoTemplateConfig.from_dict(config["model_config"]) except (TypeError, ValueError) as exc: raise HubLoadError(f"invalid model_config: {exc}") from exc _validate_tokenizer(tokenizer, max_chains=model_config.student.max_chains) return ReleaseSpec( root=root, config=MappingProxyType(config), tokenizer=MappingProxyType(tokenizer), model_config=model_config, precision="bfloat16", ) def _resolve_runtime_dtype(torch: Any, value: str | None) -> Any: name = "bfloat16" if value in (None, "auto") else value mapping = { "bfloat16": torch.bfloat16, "float32": torch.float32, } if name not in mapping: raise HubLoadError("dtype must be auto, bfloat16, or float32") return mapping[name] def load_release( model_dir: str | pathlib.Path, *, device: str = "cpu", dtype: str | None = "auto", ) -> LoadedRelease: """Construct ``EvoTemplateNA`` and strictly assign BF16 safetensors. The architecture is first created on the meta device, so loading the full release never creates a second randomly initialized 1.2 GB parameter set. No training checkpoint or AlphaFold weights are read. """ spec = load_release_spec(model_dir) try: import torch from safetensors.torch import load_file from evotemplate_na.model import EvoTemplateNA except ImportError as exc: raise RuntimeError("AF3-NA+ requires torch and safetensors") from exc resolved_device = torch.device(device) if resolved_device.type == "cuda" and not torch.cuda.is_available(): raise HubLoadError("CUDA was requested but is unavailable") runtime_dtype = _resolve_runtime_dtype(torch, dtype) checkpoint_file = spec.config.get("checkpoint_file", "model.safetensors") if checkpoint_file != "model.safetensors": raise HubLoadError("checkpoint_file must be model.safetensors") weights_path = spec.root / checkpoint_file if not weights_path.is_file(): raise HubLoadError("model.safetensors is missing") model = EvoTemplateNA( spec.model_config, device=torch.device("meta"), dtype=torch.bfloat16 ) expected = model.state_dict() declared_parameter_count = spec.config.get("parameter_count") if declared_parameter_count is not None and declared_parameter_count != sum( int(value.numel()) for value in expected.values() ): raise HubLoadError("parameter_count differs from the configured architecture") try: state = load_file(str(weights_path), device="cpu") except Exception as exc: raise HubLoadError(f"cannot load model.safetensors: {exc}") from exc expected_keys = set(expected) actual_keys = set(state) if actual_keys != expected_keys: raise HubLoadError( "safetensors key set differs from the architecture: " f"missing={sorted(expected_keys - actual_keys)} " f"extra={sorted(actual_keys - expected_keys)}" ) for name in sorted(expected_keys): value = state[name] if tuple(value.shape) != tuple(expected[name].shape): raise HubLoadError( f"tensor {name!r} shape {tuple(value.shape)} != " f"{tuple(expected[name].shape)}" ) if value.dtype != torch.bfloat16: raise HubLoadError( f"tensor {name!r} has dtype {value.dtype}; expected torch.bfloat16" ) try: incompatible = model.load_state_dict(state, strict=True, assign=True) except (RuntimeError, TypeError) as exc: raise HubLoadError(f"strict safetensors assignment failed: {exc}") from exc if incompatible.missing_keys or incompatible.unexpected_keys: raise HubLoadError("strict safetensors assignment returned incompatible keys") del state, expected try: model = model.to(device=resolved_device, dtype=runtime_dtype) except RuntimeError as exc: raise HubLoadError(f"cannot place model on {resolved_device}: {exc}") from exc model.requires_grad_(False) model.eval() return LoadedRelease( model=model, spec=spec, device=resolved_device, dtype=runtime_dtype, ) __all__ = [ "EXPECTED_VOCABULARY", "HubLoadError", "LoadedRelease", "ReleaseSpec", "load_release", "load_release_spec", ]