| """Reloadable inference artifact export.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from collections.abc import Mapping |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
|
|
| from graph.config import GraphConfig, PATConfig |
| from graph.model import GraphModel |
|
|
|
|
| def load_checkpoint_file( |
| path: str | Path, *, map_location: str | torch.device = "cpu" |
| ) -> dict[str, Any]: |
| """Load a GRM checkpoint with ``weights_only=True``. |
| |
| Checkpoints written by :class:`~graph.training.trainer.GraphTrainer` contain |
| only tensors and plain containers, so unpickling arbitrary Python objects is |
| never required. Loading model artifacts with ``weights_only=False`` would |
| permit code execution from a checkpoint file, which matters because the |
| registry promotion flow (PLAN section 40) moves artifacts between machines. |
| """ |
| |
| |
| payload = torch.load(path, map_location=map_location, weights_only=True) |
| if not isinstance(payload, dict): |
| raise TypeError(f"{path}: checkpoint is not a dict payload") |
| return payload |
|
|
|
|
| def save_inference_artifact( |
| model: GraphModel, directory: str | Path, *, use_safetensors: bool = False |
| ) -> Path: |
| root = Path(directory) |
| root.mkdir(parents=True, exist_ok=True) |
| if use_safetensors: |
| try: |
| from safetensors.torch import save_file |
| except ( |
| ImportError |
| ) as exc: |
| raise RuntimeError( |
| "safetensors export requested but safetensors is not installed" |
| ) from exc |
| weights = root / "model.safetensors" |
| save_file( |
| {key: value.detach().cpu() for key, value in model.state_dict().items()}, |
| str(weights), |
| ) |
| weight_format = "safetensors" |
| else: |
| weights = root / "model.pt" |
| torch.save(model.state_dict(), weights) |
| weight_format = "torch" |
| config = model.config.to_dict() |
| config["weight_format"] = weight_format |
| (root / "config.json").write_text( |
| json.dumps(config, indent=2, sort_keys=True), encoding="utf-8" |
| ) |
| return weights |
|
|
|
|
| def load_inference_artifact( |
| directory: str | Path, *, map_location: str | torch.device = "cpu" |
| ) -> GraphModel: |
| root = Path(directory) |
| raw = json.loads((root / "config.json").read_text(encoding="utf-8")) |
| weight_format = raw.pop("weight_format", "torch") |
| pat = PATConfig(**raw.pop("pat")) |
| raw.pop("training", None) |
| raw.pop("loss", None) |
| config = GraphConfig(pat=pat, **raw) |
| model = GraphModel(config) |
| if weight_format == "safetensors": |
| try: |
| from safetensors.torch import load_file |
| except ( |
| ImportError |
| ) as exc: |
| raise RuntimeError( |
| "safetensors artifact requires the safetensors package" |
| ) from exc |
| state = load_file(str(root / "model.safetensors"), device=str(map_location)) |
| else: |
| state = torch.load( |
| root / "model.pt", map_location=map_location, weights_only=True |
| ) |
| load_state_dict_compat(model, state) |
| model.eval() |
| return model |
|
|
|
|
| def load_state_dict_compat( |
| model: GraphModel, state: Mapping[str, torch.Tensor] |
| ) -> None: |
| """Load checkpoints across explicitly supported PAT shape migrations. |
| |
| The structured PAT token expansion added type/attribute tables and a |
| larger position table. Older registered smoke artifacts remain |
| loadable by zero-initializing only those new inputs and padding the old |
| position rows; every other missing or incompatible tensor fails loudly. |
| """ |
| expected = model.state_dict() |
| migrated: dict[str, torch.Tensor] = {} |
| missing_allowed = { |
| "pat.type_embedding.weight", |
| "pat.attribute_embedding.weight", |
| } |
| for key, value in state.items(): |
| if key not in expected: |
| raise RuntimeError(f"checkpoint contains unexpected tensor: {key}") |
| if value.shape == expected[key].shape: |
| migrated[key] = value |
| continue |
| if ( |
| key == "pat.path_position_embedding.weight" |
| and value.ndim == 2 |
| and value.shape[1] == expected[key].shape[1] |
| and value.shape[0] < expected[key].shape[0] |
| ): |
| padded = torch.zeros_like(expected[key]) |
| padded[: value.shape[0]] = value |
| migrated[key] = padded |
| continue |
| raise RuntimeError( |
| f"checkpoint tensor shape mismatch for {key}: {tuple(value.shape)} != {tuple(expected[key].shape)}" |
| ) |
| for key, value in expected.items(): |
| if key not in migrated: |
| if key in missing_allowed: |
| migrated[key] = torch.zeros_like(value) |
| else: |
| raise RuntimeError(f"checkpoint is missing tensor: {key}") |
| model.load_state_dict(migrated, strict=True) |
|
|
