torch_io.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Trusted local checkpoint loads (PyTorch 2.6+ weights_only default)."""
from __future__ import annotations
 
import json
from pathlib import Path
from typing import Any
 
 
def is_corrupt_pytorch_zip(exc: BaseException) -> bool:
    msg = str(exc).lower()
    return (
        "zip archive" in msg
        or "central directory" in msg
        or "pytorchstreamreader" in msg
    )
 
 
def load_checkpoint(path: str | Path, *, map_location: Any = None) -> Any:
    """Load a trusted Aquin .pt checkpoint (SAE, norm stats, activations)."""
    import torch
 
    kwargs: dict[str, Any] = {}
    if map_location is not None:
        kwargs["map_location"] = map_location
    try:
        try:
            return torch.load(str(path), weights_only=False, **kwargs)
        except TypeError:
            return torch.load(str(path), **kwargs)
    except Exception as exc:
        if is_corrupt_pytorch_zip(exc):
            raise ValueError(
                f"Checkpoint at {path} is corrupted or incomplete "
                "(download may have been interrupted). Delete the file and download again."
            ) from exc
        raise
 
 
def load_norm_stats(path: str | Path, *, map_location: Any = None) -> dict[str, Any] | None:
    """Load {mean, std} norm stats from .pt (torch) or .json sidecar."""
    import torch
 
    p = Path(path)
    if not p.is_file():
        return None
 
    if p.suffix.lower() == ".json" or p.read_bytes()[:1] == b"{":
        if looks_like_catalog_metadata_blob(p):
            print(
                f"[norm] {p} is catalog metadata, not norm stats — "
                f"remove it and re-run: aquin load sae <model>-l<layer>",
                flush=True,
            )
            return None
        try:
            data = json.loads(p.read_text(encoding="utf-8"))
        except Exception:
            return None
        if not isinstance(data, dict) or "mean" not in data or "std" not in data:
            return None
        mean = data["mean"]
        std = data["std"]
        if not isinstance(mean, torch.Tensor):
            mean = torch.tensor(mean)
        if not isinstance(std, torch.Tensor):
            std = torch.tensor(std)
        if map_location is not None:
            mean = mean.to(map_location)
            std = std.to(map_location)
        return {"mean": mean, "std": std}
 
    try:
        data = load_checkpoint(p, map_location=map_location)
    except Exception as exc:
        print(f"[norm] could not load {p}: {exc} — continuing without normalization", flush=True)
        return None
 
    if looks_like_catalog_metadata_dict(data):
        print(
            f"[norm] {p} is catalog metadata, not norm stats — "
            f"remove it and re-run: aquin load sae <model>-l<layer>",
            flush=True,
        )
        return None
 
    if not isinstance(data, dict) or "mean" not in data or "std" not in data:
        print(f"[norm] unexpected format in {p} — continuing without normalization", flush=True)
        return None
 
    mean, std = data["mean"], data["std"]
    if map_location is not None and hasattr(mean, "to"):
        mean = mean.to(map_location)
        std = std.to(map_location)
    return {"mean": mean, "std": std}
 
 
def looks_like_catalog_metadata_blob(path: str | Path) -> bool:
    """True when a .pt file is actually public_saes catalog JSON (not norm stats)."""
    p = Path(path)
    if not p.is_file() or p.stat().st_size < 2 or p.read_bytes()[:1] != b"{":
        return False
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
    except Exception:
        return False
    if not isinstance(data, dict):
        return False
    if "mean" in data and "std" in data:
        return False
    return "model_slug" in data or "d_model" in data or "d_sae" in data
 
 
def looks_like_invalid_norm_blob(path: str | Path) -> bool:
    """True when a norm file is catalog JSON or other non-norm JSON."""
    if looks_like_catalog_metadata_blob(path):
        return True
    p = Path(path)
    if not p.is_file() or p.read_bytes()[:1] != b"{":
        return False
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
    except Exception:
        return True
    return not (isinstance(data, dict) and "mean" in data and "std" in data)
 
 
def looks_like_json_error_blob(path: str | Path) -> bool:
    """True when a downloaded artifact is likely an API error body, not weights."""
    p = Path(path)
    if not p.is_file() or p.stat().st_size < 2:
        return True
    return p.read_bytes()[:1] == b"{"
 
 
def looks_like_catalog_metadata_dict(data: Any) -> bool:
    """True when a loaded checkpoint dict is public_saes row metadata, not SAE weights."""
    import torch
 
    if not isinstance(data, dict):
        return False
    if "mean" in data and "std" in data:
        return False
    if any(isinstance(v, torch.Tensor) for v in data.values()):
        return False
    if "state_dict" in data or "W_enc" in data or "W_dec" in data:
        return False
    return "model_slug" in data or ("d_model" in data and "layer" in data)
 
 
def _checkpoint_tensors(ckpt: Any) -> dict[str, Any]:
    import torch
 
    if not isinstance(ckpt, dict):
        return {}
    if "state_dict" in ckpt and isinstance(ckpt["state_dict"], dict):
        return {k: v for k, v in ckpt["state_dict"].items() if isinstance(v, torch.Tensor)}
    if "model_state" in ckpt and isinstance(ckpt["model_state"], dict):
        return {k: v for k, v in ckpt["model_state"].items() if isinstance(v, torch.Tensor)}
    return {k: v for k, v in ckpt.items() if isinstance(v, torch.Tensor)}
 
 
def _has_sae_weight_tensors(tensors: dict[str, Any]) -> bool:
    if not tensors:
        return False
    names = set(tensors.keys())
    leaves = {k.split(".")[-1] for k in names}
    sae_keys = {"W_enc", "W_dec", "encoder.weight", "decoder.weight"}
    return bool(names & sae_keys) or bool(leaves & {"W_enc", "W_dec", "weight"})
 
 
def is_valid_sae_checkpoint_path(path: str | Path) -> bool:
    """False when a cached SAE file is catalog metadata or another non-checkpoint blob."""
    p = Path(path)
    if not p.is_file() or p.stat().st_size < 64:
        return False
    if looks_like_catalog_metadata_blob(p) or looks_like_json_error_blob(p):
        return False
 
    if p.suffix.lower() == ".safetensors":
        try:
            from safetensors.torch import load_file
 
            tensors = load_file(str(p), device="cpu")
            return _has_sae_weight_tensors(tensors)
        except Exception:
            return False
 
    try:
        ckpt = load_checkpoint(p, map_location="cpu")
    except Exception:
        return False
 
    if looks_like_catalog_metadata_dict(ckpt):
        return False
    return _has_sae_weight_tensors(_checkpoint_tensors(ckpt))