Blame for src/checkpoint.py

"""
Canonical Checkpoint Management & Provenance Verification (`src/checkpoint.py`).

Enforces strict scientific integrity rules:
  - Never silently ignores missing, unexpected, or shape-mismatched checkpoint keys
  - Explicitly distinguishes trained weights (`is_trained=True`) from random initialization
  - Records full provenance metadata (dataset, initial/final loss & perplexity, seed, SHA-256, tokenizer fingerprint)
  - Fails loudly with actionable diagnostics when given an incompatible legacy artifact (e.g., `qtensorformer_small.pt`)
"""

import datetime
import hashlib
import json
import platform
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from .config import ModelConfig, validate_model_config
from .models import DenseBaseline, QTensorFormer
from .tokenizer import QTensorFormerTokenizer


def compute_file_sha256(filepath: Union[str, Path]) -> str:
    path = Path(filepath)
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def compute_state_dict_fingerprint(state_dict: Dict[str, torch.Tensor]) -> str:
    h = hashlib.sha256()
    for key in sorted(state_dict.keys()):
        tensor = state_dict[key].detach().cpu().float().contiguous()
        h.update(key.encode("utf-8"))
        h.update(str(tuple(tensor.shape)).encode("utf-8"))
        h.update(tensor.numpy().tobytes())
    return h.hexdigest()


def inspect_checkpoint_compatibility(
    checkpoint_path: Union[str, Path],
    model: nn.Module,
) -> Dict[str, Any]:
    """
    Forensically inspect a checkpoint file against a target `nn.Module` instance.
    """
    path = Path(checkpoint_path)
    if not path.exists():
        raise FileNotFoundError(f"Checkpoint file does not exist: {path}")

    raw = torch.load(path, map_location="cpu", weights_only=False)
    if isinstance(raw, dict) and "model_state_dict" in raw:
        ckpt_sd = raw["model_state_dict"]
    elif isinstance(raw, dict) and "state_dict" in raw:
        ckpt_sd = raw["state_dict"]
    elif isinstance(raw, dict):
        ckpt_sd = raw
    else:
        raise ValueError(f"Unsupported checkpoint container type in {path}: {type(raw).__name__}")

    model_sd = model.state_dict()
    ckpt_keys = set(ckpt_sd.keys())
    model_keys = set(model_sd.keys())

    matched_keys = sorted(ckpt_keys & model_keys)
    missing_keys = sorted(model_keys - ckpt_keys)
    unexpected_keys = sorted(ckpt_keys - model_keys)

    shape_mismatches = []
    has_nan_or_inf = False
    for k in matched_keys:
        c_shape = tuple(ckpt_sd[k].shape)
        m_shape = tuple(model_sd[k].shape)
        if c_shape != m_shape:
            shape_mismatches.append({"key": k, "checkpoint_shape": c_shape, "model_shape": m_shape})
    for k, v in ckpt_sd.items():
        if isinstance(v, torch.Tensor) and v.is_floating_point():
            if not torch.isfinite(v).all():
                has_nan_or_inf = True
                break

    strict_compatible = (
        len(missing_keys) == 0
        and len(unexpected_keys) == 0
        and len(shape_mismatches) == 0
        and not has_nan_or_inf
    )

    return {
        "checkpoint_path": str(path),
        "file_size_bytes": path.stat().st_size,
        "file_sha256": compute_file_sha256(path),
        "target_model_class": model.__class__.__name__,
        "checkpoint_key_count": len(ckpt_keys),
        "model_key_count": len(model_keys),
        "matched_key_count": len(matched_keys),
        "missing_key_count": len(missing_keys),
        "unexpected_key_count": len(unexpected_keys),
        "shape_mismatch_count": len(shape_mismatches),
        "shape_mismatches": shape_mismatches[:10],
        "sample_missing_keys": missing_keys[:10],
        "sample_unexpected_keys": unexpected_keys[:10],
        "has_nan_or_inf": has_nan_or_inf,
        "strict_compatible": strict_compatible,
        "is_trained": bool(raw.get("is_trained", False)) if isinstance(raw, dict) else False,
        "architecture": raw.get("architecture", "unknown") if isinstance(raw, dict) else "unknown",
        "training_provenance": raw.get("training_provenance", None) if isinstance(raw, dict) else None,
    }


def save_canonical_checkpoint(
    model: Union[QTensorFormer, DenseBaseline, nn.Module],
    checkpoint_path: Union[str, Path],
    tokenizer: Optional[QTensorFormerTokenizer] = None,
    is_trained: bool = False,
    training_provenance: Optional[Dict[str, Any]] = None,
    metadata_path: Optional[Union[str, Path]] = None,
) -> Dict[str, Any]:
    """
    Save a canonical Q-TensorFormer or DenseBaseline checkpoint with full provenance metadata.
    """
    path = Path(checkpoint_path)
    path.parent.mkdir(parents=True, exist_ok=True)

    state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}
    for k, v in state_dict.items():
        if v.is_floating_point() and not torch.isfinite(v).all():
            raise ValueError(f"Refusing to save checkpoint: parameter '{k}' contains NaN or Inf values.")

    cfg_dict = {}
    if hasattr(model, "config") and model.config is not None:
        if hasattr(model.config, "__dict__"):
            cfg_dict = {k: v for k, v in model.config.__dict__.items() if not k.startswith("_")}

    tok_meta = None
    if tokenizer is not None:
        tokenizer.validate_model_alignment(model)
        tok_meta = {
            "tokenizer_class": tokenizer.__class__.__name__,
            "version": tokenizer.FORMAT_VERSION,
            "vocab_size": tokenizer.vocab_size,
            "vocab_sha256": tokenizer.fingerprint(),
            "special_tokens": {
                "pad_token_id": tokenizer.pad_token_id,
                "bos_token_id": tokenizer.bos_token_id,
                "eos_token_id": tokenizer.eos_token_id,
                "unk_token_id": tokenizer.unk_token_id,
            },
        }

    sd_sha256 = compute_state_dict_fingerprint(state_dict)
    created_utc = datetime.datetime.now(datetime.timezone.utc).isoformat()

    payload = {
        "format": "Q-TensorFormer-Canonical-Checkpoint-v1",
        "architecture": model.__class__.__name__,
        "created_utc": created_utc,
        "is_trained": bool(is_trained),
        "weight_status": "TRAINED_CANONICAL_WEIGHTS" if is_trained else "RANDOM_INITIALIZATION",
        "model_config": cfg_dict,
        "tokenizer_metadata": tok_meta,
        "training_provenance": training_provenance or {},
        "state_dict_sha256": sd_sha256,
        "parameter_count": sum(p.numel() for p in model.parameters()),
        "state_dict_key_count": len(state_dict),
        "software": {
            "python": platform.python_version(),
            "pytorch": torch.__version__,
            "platform": platform.platform(),
        },
        "model_state_dict": state_dict,
    }

    torch.save(payload, path)
    file_sha256 = compute_file_sha256(path)

    meta_summary = {k: v for k, v in payload.items() if k != "model_state_dict"}
    meta_summary["checkpoint_path"] = str(path)
    meta_summary["file_size_bytes"] = path.stat().st_size
    meta_summary["file_sha256"] = file_sha256

    if metadata_path is None:
        metadata_path = path.with_suffix(".metadata.json")
    else:
        metadata_path = Path(metadata_path)
    with open(metadata_path, "w", encoding="utf-8") as f:
        json.dump(meta_summary, f, indent=2)

    return meta_summary


def load_canonical_checkpoint(
    checkpoint_path: Union[str, Path],
    model: Optional[nn.Module] = None,
    strict: bool = True,
    require_trained: bool = False,
    tokenizer: Optional[QTensorFormerTokenizer] = None,
    map_location: Union[str, torch.device] = "cpu",
) -> Tuple[nn.Module, Dict[str, Any]]:
    """
    Load and validate a canonical checkpoint with strict key/shape/provenance verification.
    Never silently falls back to random weights or ignores missing keys.
    """
    path = Path(checkpoint_path)
    if not path.exists():
        raise FileNotFoundError(f"Checkpoint file not found: {path}")

    raw = torch.load(path, map_location=map_location, weights_only=False)
    if not isinstance(raw, dict):
        raise ValueError(f"Invalid checkpoint container in {path}: expected dict, got {type(raw).__name__}")

    if "model_state_dict" in raw:
        state_dict = raw["model_state_dict"]
    elif "state_dict" in raw:
        state_dict = raw["state_dict"]
    else:
        state_dict = raw

    if model is None:
        cfg_data = raw.get("model_config") or raw.get("config")
        if not isinstance(cfg_data, dict):
            raise ValueError(
                f"Cannot instantiate model from {path}: checkpoint does not contain a valid 'model_config' dict."
            )
        valid_fields = set(ModelConfig.__dataclass_fields__.keys())
        filtered_cfg = {k: v for k, v in cfg_data.items() if k in valid_fields}
        cfg = ModelConfig(**filtered_cfg)
        validate_model_config(cfg, QTensorFormer)
        arch = raw.get("architecture", "QTensorFormer")
        if arch == "DenseBaseline":
            model = DenseBaseline(cfg)
        else:
            model = QTensorFormer(cfg)

    compat = inspect_checkpoint_compatibility(path, model)
    if strict and not compat["strict_compatible"]:
        raise RuntimeError(
            f"CheckpointIncompatibilityError: Checkpoint '{path.name}' is incompatible with "
            f"{model.__class__.__name__} under strict=True.\n"
            f"  - Matched keys: {compat['matched_key_count']}/{compat['model_key_count']}\n"
            f"  - Missing keys ({compat['missing_key_count']}): {compat['sample_missing_keys'][:5]}\n"
            f"  - Unexpected keys ({compat['unexpected_key_count']}): {compat['sample_unexpected_keys'][:5]}\n"
            f"  - Shape mismatches ({compat['shape_mismatch_count']}): {compat['shape_mismatches'][:3]}\n"
            f"  - Has NaN/Inf: {compat['has_nan_or_inf']}\n"
            "Use the canonical trained checkpoint 'checkpoints/qtensorformer_canonical_trained.pt' instead."
        )

    if require_trained and not compat["is_trained"]:
        raise RuntimeError(
            f"UntrainedCheckpointError: Checkpoint '{path.name}' does not have verified training provenance "
            "(`is_trained=True`). Refusing to treat untrained weights as a trained checkpoint."
        )

    if tokenizer is not None:
        tokenizer.validate_model_alignment(model)
        tok_meta = raw.get("tokenizer_metadata")
        if isinstance(tok_meta, dict) and "vocab_sha256" in tok_meta:
            if tok_meta["vocab_sha256"] != tokenizer.fingerprint():
                raise ValueError(
                    f"TokenizerChecksumMismatchError: Checkpoint expects tokenizer SHA-256 "
                    f"{tok_meta['vocab_sha256']}, but provided tokenizer has {tokenizer.fingerprint()}."
                )

    model.load_state_dict(state_dict, strict=strict)
    model.to(map_location)
    return model, compat