giant/giant/checkpoint_io.py

"""Load a trained checkpoint into ready-to-run models (giant.cli's `predict`/`rollout`).

Both commands need the same ~15 steps to go from a checkpoint path to two
`eval()`-mode models plus their normalizers/vocab maps: load the pickle,
validate it carries what current code expects, resolve which conditioning
mode each axis was trained with, restore the top-N vocab maps (if the
checkpoint used one-hot conditioning), rebuild the normalizers, construct the
model from `model_config`, and load the requested (raw or EMA) weights. This
used to be duplicated near-verbatim in both commands (issues.md Issue 5) —
`load_for_inference` is the single implementation.

This module intentionally has no Typer dependency, so it can be unit-tested
directly and imported from non-CLI code (`giant.analysis.router_gating`,
lazily — see that module's docstring for why). Failures raise
`CheckpointCompatibilityError` with the same wording the CLI has always
shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import torch
from torch import nn

from giant import config as gconfig
from giant.constants import K_MAX
from giant.data.loader import TopNMap
from giant.data.setup_cache import topnmap_from_json
from giant.data.transforms import Normalizer
from giant.model.network import build_models


class CheckpointCompatibilityError(Exception):
    """Checkpoint is missing something `load_for_inference` needs."""


def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
    """(particle_conditioning, material_conditioning) for
    `giant.data.transforms.build_cond_features`/`build_features` — from
    either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
    shared string, same for both axes) or a new-format one (independent
    `model_config["conditioning"]["particle"/"material"]["type"]` — the two
    axes are configured independently and may differ)."""
    raw = model_cfg.get("conditioning", default)
    if isinstance(raw, dict):
        return (
            raw.get("particle", {}).get("type", default),
            raw.get("material", {}).get("type", default),
        )
    return raw, raw


def stage_cfg(model_cfg: dict, stage: str) -> dict:
    """`model_cfg[f"{stage}_model"]` for a new-format model_config, `{}` for
    a v0.2 flat one (whose ddpm schedule always used `CosineSchedule`'s own
    default `T=1000` — never a config key — and which never had
    `particle_type` at all, so `{}` is the correct fallback for both
    `ddpm_steps`/`particle_type_other_policy` below)."""
    val = model_cfg.get(f"{stage}_model")
    return val if isinstance(val, dict) else {}


def ddpm_steps(model_cfg: dict, stage: str) -> int:
    return stage_cfg(model_cfg, stage).get("ddpm", {}).get("n_steps", 1000)


def particle_type_other_policy(model_cfg: dict) -> str:
    return stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")


def load_pdg_topn_map(ckpt: dict) -> TopNMap | None:
    """`ckpt["pdg_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
    this checkpoint's conditioning/particle_type never needed one (see
    `giant.pipeline.run_setup_stage`, which only populates it when
    `conditioning.particle.type` or `stage2_model.particle_type.target` is
    `"onehot"`)."""
    raw = ckpt.get("pdg_topn_map")
    return topnmap_from_json(raw, axis="pdg") if raw is not None else None


def load_mat_topn_map(ckpt: dict) -> TopNMap | None:
    """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
    this checkpoint's `conditioning.material.type` was never `"onehot"` (see
    `giant.pipeline.run_setup_stage`)."""
    raw = ckpt.get("mat_topn_map")
    return topnmap_from_json(raw, axis="material") if raw is not None else None


def load_sec_type_topn_map(ckpt: dict) -> TopNMap | None:
    """`ckpt["sec_type_topn_map"]` as a `giant.data.loader.TopNMap`, or
    `None` if this checkpoint's `stage2_model.particle_type.target` was never
    `"onehot"` (see `giant.pipeline.run_setup_stage`).

    Pre-gitea-#29 checkpoints have no `sec_type_topn_map` key at all — before
    #29, the secondary-species decode map and the conditioning PDG onehot map
    were always numerically the same map, saved once under `pdg_topn_map`.
    For those, fall back to `load_pdg_topn_map` to reproduce that exact
    behavior; a current checkpoint always has the key (possibly `null`, if
    `particle_type.target != "onehot"`), so this fallback never fires for one."""
    if "sec_type_topn_map" in ckpt:
        raw = ckpt["sec_type_topn_map"]
        return topnmap_from_json(raw, axis="pdg") if raw is not None else None
    return load_pdg_topn_map(ckpt)


@dataclass(frozen=True)
class InferenceContext:
    """Everything needed to run a trained checkpoint forward, resolved once."""

    stage1: nn.Module | None
    stage2: nn.Module | None
    cond_norm: Normalizer
    tgt_norm: Normalizer
    sec_phys_norm: Normalizer
    pdg_map: dict[int, int]
    mat_map: dict[str, int]
    pdg_topn_map: TopNMap | None
    mat_topn_map: TopNMap | None
    sec_type_topn_map: TopNMap | None
    particle_conditioning: str
    material_conditioning: str
    k_max: int
    stage1_ddpm_steps: int
    stage2_ddpm_steps: int
    other_policy: str
    model_config: dict
    epoch: int | None
    best_val_loss: float | None


def load_for_inference(
    checkpoint: Path,
    device: torch.device,
    command_name: str,
    weights: str = "raw",
    require_stage2: bool = True,
) -> InferenceContext:
    """Load *checkpoint* and reconstruct everything `predict`/`rollout` need
    to run it forward, on *device*, in `eval()` mode.

    *command_name* (e.g. `"predict"`/`"rollout"`) only feeds the "needs both"
    error message below. *weights* is `"raw"` (the live training weights) or
    `"ema"` (the EMA shadow copy, see `--ema-decay`). *require_stage2*
    controls whether a checkpoint with an inactive stage 2
    (`stage2_model.active = false`) is an error (both current callers need
    both stages) or an acceptable `stage2 = None` result — kept as a real
    parameter since `stage{1,2}_model.active` is a real, if currently
    stage1+stage2-only-in-practice, config option.
    """
    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
    for key in ("model_config", "sec_decoder"):
        if key not in ckpt:
            raise CheckpointCompatibilityError(f"checkpoint has no {key} — retrain with the current code")

    if "sec_phys" not in ckpt.get("normalizer", {}):
        raise CheckpointCompatibilityError("checkpoint has no normalizer.sec_phys — retrain with the current code")

    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)

    model_cfg = ckpt["model_config"]
    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
    pdg_topn_map = load_pdg_topn_map(ckpt)
    mat_topn_map = load_mat_topn_map(ckpt)
    if particle_conditioning == "onehot" and pdg_topn_map is None:
        raise CheckpointCompatibilityError(
            "checkpoint's conditioning.particle.type='onehot' but has no pdg_topn_map — retrain with the current code"
        )
    if material_conditioning == "onehot" and mat_topn_map is None:
        raise CheckpointCompatibilityError(
            "checkpoint's conditioning.material.type='onehot' but has no mat_topn_map — retrain with the current code"
        )
    sec_type_topn_map = load_sec_type_topn_map(ckpt)
    particle_type_target = stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("target", "onehot")
    if particle_type_target == "onehot" and sec_type_topn_map is None:
        raise CheckpointCompatibilityError(
            "checkpoint's stage2_model.particle_type.target='onehot' but has no "
            "sec_type_topn_map — retrain with the current code"
        )
    other_policy = particle_type_other_policy(model_cfg)
    stage1_ddpm_steps = ddpm_steps(model_cfg, "stage1")
    stage2_ddpm_steps = ddpm_steps(model_cfg, "stage2")
    k_max = stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])

    built = build_models(model_cfg)
    stage1, stage2 = built["stage1"], built["stage2"]
    if require_stage2 and (stage1 is None or stage2 is None):
        raise CheckpointCompatibilityError(
            f"checkpoint has an inactive stage1 or stage2 — {command_name} needs both (see stage{{1,2}}_model.active)"
        )

    if weights == "raw":
        model_key, sec_key = "model", "sec_decoder"
    else:
        model_key, sec_key = "model_ema", "sec_decoder_ema"
        if model_key not in ckpt or sec_key not in ckpt:
            raise CheckpointCompatibilityError(
                f"{checkpoint} has no EMA weights (trained before --ema-decay, "
                "or with --ema-decay 0) — use --weights raw"
            )
    if stage1 is not None:
        stage1.load_state_dict(ckpt[model_key])
        stage1.to(device).eval()
    if stage2 is not None:
        stage2.load_state_dict(ckpt[sec_key])
        stage2.to(device).eval()

    return InferenceContext(
        stage1=stage1,
        stage2=stage2,
        cond_norm=cond_norm,
        tgt_norm=tgt_norm,
        sec_phys_norm=sec_phys_norm,
        pdg_map=pdg_map,
        mat_map=mat_map,
        pdg_topn_map=pdg_topn_map,
        mat_topn_map=mat_topn_map,
        sec_type_topn_map=sec_type_topn_map,
        particle_conditioning=particle_conditioning,
        material_conditioning=material_conditioning,
        k_max=k_max,
        stage1_ddpm_steps=stage1_ddpm_steps,
        stage2_ddpm_steps=stage2_ddpm_steps,
        other_policy=other_policy,
        model_config=model_cfg,
        epoch=ckpt.get("epoch"),
        best_val_loss=ckpt.get("best_val_loss"),
    )