| """Load a published checkpoint. |
| |
| Each published run is a directory holding |
| |
| model.safetensors the parameters |
| config.json {"kind", "cfg", "train_args", "params", ...} |
| |
| `kind` selects the architecture; `cfg` is the exact dataclass field set the run |
| was built with, so nothing is inferred from tensor shapes. |
| """ |
| from __future__ import annotations |
|
|
| import json |
|
|
| import torch |
| from safetensors.torch import load_file |
|
|
| from .common import ckpt_root |
|
|
|
|
| def load_config(run: str) -> dict: |
| return json.loads((ckpt_root() / run / "config.json").read_text()) |
|
|
|
|
| def build(kind: str, cfg: dict): |
| if kind == "rmt": |
| from .model.rmt import RMT, RMTConfig |
| return RMT(RMTConfig(**{k: v for k, v in cfg.items() |
| if k in RMTConfig.__dataclass_fields__})) |
| if kind in ("pnsr", "pnsr_k1"): |
| from .model.pnsr import PNSR, PNSRConfig |
| return PNSR(PNSRConfig(**{k: v for k, v in cfg.items() |
| if k in PNSRConfig.__dataclass_fields__})) |
| if kind in ("bind", "unbound"): |
| from .model.bind import BindConfig, PNSBind |
| return PNSBind(BindConfig(**{k: v for k, v in cfg.items() |
| if k in BindConfig.__dataclass_fields__})) |
| from .model.tx import TX, TXConfig |
| return TX(TXConfig(**{k: v for k, v in cfg.items() |
| if k in TXConfig.__dataclass_fields__})) |
|
|
|
|
| def load_model(run: str, device: str = "cuda"): |
| """Returns (model in eval mode on `device`, kind, config dict).""" |
| conf = load_config(run) |
| model = build(conf["kind"], conf["cfg"]) |
| state = load_file(str(ckpt_root() / run / "model.safetensors"), device="cpu") |
| model.load_state_dict(state, strict=True) |
| model.eval().to(device) |
| return model, conf["kind"], conf |
|
|
|
|
| def param_count(model) -> int: |
| """Number of distinct trainable parameters — the paper's budget figure. |
| |
| The event encoder and the record encoder *share* one 8192x384 token |
| embedding table (the parameter is tied, not duplicated), so `state_dict` |
| lists it under two keys while the model holds it once. `parameters()` |
| already de-duplicates shared tensors, which is exactly the quantity the |
| matched-parameter comparison uses. |
| """ |
| return sum(p.numel() for p in model.parameters()) |
|
|
|
|
| @torch.no_grad() |
| def verify(run: str, device: str = "cpu") -> dict: |
| """Load a checkpoint and check it against its recorded provenance.""" |
| conf = load_config(run) |
| model, _, _ = load_model(run, device) |
| n = param_count(model) |
| tensors = sum(v.numel() for v in model.state_dict().values()) |
| return dict(run=run, kind=conf["kind"], loaded=True, |
| params=n, params_match=n == conf["params"], |
| state_dict_tensor_total=tensors, |
| tensor_total_match=tensors == conf["state_dict_tensor_total"]) |
|
|
