inspect_checkpoint.py

"""
Inspect a training checkpoint and write a cleaned inference-only state.
- Drops optimizer variables.
- If ema_trainable exists, keeps only ema_trainable + non_trainable (drops trainable).
- Saves to weights/gen_state.pkl for use by inference.py / app.py.
"""
import os
import sys
import pickle

def _describe(obj, depth=0, max_depth=3):
    """Return a short description of the object for printing."""
    if depth > max_depth:
        return "..."
    if hasattr(obj, "shape"):
        return f"array shape={obj.shape} dtype={getattr(obj, 'dtype', '?')}"
    if isinstance(obj, dict):
        return "dict(" + ", ".join(f"{k!r}" for k in list(obj.keys())[:8]) + ("..." if len(obj) > 8 else "") + ")"
    if isinstance(obj, (list, tuple)):
        return f"{type(obj).__name__} len={len(obj)}"
    return type(obj).__name__


def inspect_checkpoint(path):
    """Load checkpoint and print top-level keys and a short description of each value."""
    with open(path, "rb") as f:
        data = pickle.load(f)

    if not isinstance(data, dict):
        print(f"Top-level type: {type(data)}")
        return data

    print("Checkpoint keys and value summary:")
    print("-" * 60)
    for k in sorted(data.keys()):
        v = data[k]
        desc = _describe(v)
        print(f"  {k!r}: {desc}")
    print("-" * 60)
    return data


# Keys we consider optimizer / training-only (to drop)
OPT_OR_TRAINING_KEYS = {
    "opt_state", "optimizer", "optimizer_state", "opt",
    "step", "steps", "epoch", "epochs",
    "trainable",  # dropped when ema_trainable exists
    "rng",  # inference creates its own; optional to keep for reproducibility
}


def clean_checkpoint(data, keep_rng=False):
    """
    Build a state dict for inference only.
    - Drop any key in OPT_OR_TRAINING_KEYS (and 'trainable' if ema exists).
    - If 'ema_trainable' exists, drop 'trainable'.
    - Keep only ema_trainable (or trainable) + non_trainable.
    """
    has_ema = "ema_trainable" in data
    out = {}

    if has_ema:
        out["ema_trainable"] = data["ema_trainable"]
        # drop trainable (we use ema only)
    else:
        if "trainable" in data:
            out["ema_trainable"] = data["trainable"]  # inference expects key "ema_trainable"
        # else no trainable params in checkpoint

    if "non_trainable" in data:
        out["non_trainable"] = data["non_trainable"]

    if keep_rng and "rng" in data:
        out["rng"] = data["rng"]

    return out


def main():
    import argparse
    p = argparse.ArgumentParser(description="Inspect checkpoint and optionally write cleaned gen_state.pkl")
    p.add_argument("checkpoint", nargs="?", default="weights/checkpoint.pkl",
                   help="Path to full training checkpoint (default: weights/checkpoint.pkl)")
    p.add_argument("-o", "--output", default="weights/gen_state.pkl",
                   help="Output path for cleaned state (default: weights/gen_state.pkl)")
    p.add_argument("--no-write", action="store_true", help="Only inspect, do not write cleaned state")
    p.add_argument("--keep-rng", action="store_true", help="Keep rng in cleaned state")
    args = p.parse_args()

    if not os.path.isfile(args.checkpoint):
        print(f"Checkpoint not found: {args.checkpoint}", file=sys.stderr)
        sys.exit(1)

    data = inspect_checkpoint(args.checkpoint)

    if not isinstance(data, dict):
        print("Checkpoint is not a dict; cannot clean.", file=sys.stderr)
        sys.exit(1)

    cleaned = clean_checkpoint(data, keep_rng=args.keep_rng)
    print("Cleaned state keys:", list(cleaned.keys()))

    if not args.no_write:
        os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
        with open(args.output, "wb") as f:
            pickle.dump(cleaned, f)
        print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()