File size: 8,831 Bytes
a371c7e | 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 245 246 247 248 249 250 251 252 253 254 | """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",
]
|