hub_loader.py

"""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",
]