model/checkpoint.py

from __future__ import annotations

from contextlib import contextmanager
from pathlib import Path
import sys
from typing import Any, Iterator, Mapping

import torch
from torch import nn


@contextmanager
def _temporary_module_path(source_root: str | Path | None) -> Iterator[None]:
    if source_root is None:
        yield
        return

    path = str(Path(source_root).resolve())
    sys.path.insert(0, path)
    try:
        yield
    finally:
        try:
            sys.path.remove(path)
        except ValueError:
            pass


def _extract_state_dict(checkpoint: Any) -> Mapping[str, torch.Tensor]:
    if isinstance(checkpoint, Mapping):
        for key in ("model", "model_state", "state_dict"):
            candidate = checkpoint.get(key)
            if isinstance(candidate, Mapping):
                return candidate
        if checkpoint and all(isinstance(key, str) for key in checkpoint):
            return checkpoint
    raise ValueError("Checkpoint does not contain a model state dictionary.")


def load_checkpoint(
    model: nn.Module,
    checkpoint_path: str | Path,
    *,
    strict: bool = True,
    map_location: str | torch.device = "cpu",
    source_root: str | Path | None = None,
) -> dict[str, Any]:
    """Load a trusted W-MAE checkpoint and return loading metadata.

    ``source_root`` is needed only for legacy checkpoints whose pickle payload
    references modules from the original repository, such as ``utils.YParams``.
    """
    checkpoint_path = Path(checkpoint_path)
    if not checkpoint_path.is_file():
        raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint_path}")

    with _temporary_module_path(source_root):
        checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False)

    state_dict = _extract_state_dict(checkpoint)
    model_state = model.state_dict()
    missing = sorted(set(model_state) - set(state_dict))
    unexpected = sorted(set(state_dict) - set(model_state))
    shape_mismatches = {
        key: {"model": list(model_state[key].shape), "checkpoint": list(state_dict[key].shape)}
        for key in model_state.keys() & state_dict.keys()
        if tuple(model_state[key].shape) != tuple(state_dict[key].shape)
    }
    if strict and (missing or unexpected or shape_mismatches):
        raise RuntimeError(
            "Checkpoint is incompatible with the model: "
            f"missing={missing}, unexpected={unexpected}, shape_mismatches={shape_mismatches}"
        )

    incompatible = model.load_state_dict(state_dict, strict=strict)
    return {
        "checkpoint_path": str(checkpoint_path),
        "strict": strict,
        "missing_keys": list(incompatible.missing_keys),
        "unexpected_keys": list(incompatible.unexpected_keys),
        "shape_mismatches": shape_mismatches,
        "checkpoint_keys": len(state_dict),
    }