scene-actor-extraction/scripts/sae_stamp.py

"""Gallery ↔ embedder model binding (GR-004).

TRACES: GR-004 | SR-001

Python twin of src/gallery/embedder_stamp.{hpp,cpp}; the two implement the same
comparison rules and must stay in agreement. Kept as its own module — rather than
folded into sae_gallery — because scripts/optimizer/replay.py imports it once per
replay subprocess, thousands of times in a DE sweep, and must not pay for
sae_gallery's requests/Pillow imports to ask "were these made by the same model?".
Dependencies here are hashlib, json and h5py, all of which a replay already loads.

A gallery is only valid for the embedder that built it: cosine similarities across
models are meaningless but look plausible, so the mistake is silent and every
measurement taken afterwards is suspect. Identity = model filename + SHA-256 of
the model file. The hash decides (a model re-exported in place keeps its name but
not its bytes); the name is what makes the error readable. See
src/gallery/embedder_stamp.hpp for the full rationale.
"""

import hashlib
import json
import os
import sys
from pathlib import Path

import h5py


def _as_str(v) -> str:
    return v.decode() if isinstance(v, bytes) else ("" if v is None else str(v))


_STAMP_CACHE: dict = {}


class EmbedderMismatch(RuntimeError):
    """Gallery was built with a different embedder than the one about to be used."""


def sha256_file(path) -> str:
    """Lowercase hex SHA-256 of a file's bytes; "" if it cannot be read."""
    path = Path(path)
    try:
        st = path.stat()
    except OSError:
        return ""
    key = (str(path), st.st_mtime_ns, st.st_size)
    if key in _STAMP_CACHE:
        return _STAMP_CACHE[key]
    h = hashlib.sha256()
    try:
        with open(path, "rb") as f:
            for chunk in iter(lambda: f.read(1 << 20), b""):
                h.update(chunk)
    except OSError:
        return ""
    _STAMP_CACHE[key] = h.hexdigest()
    return _STAMP_CACHE[key]


def embedder_stamp(model_path, embed_dim: int = 512) -> dict:
    """Identify an embedder model file → {"model_name", "model_sha256", "embed_dim"}.

    A model file that is absent (e.g. a TRT deployment running from a prebuilt
    .engine) yields a name-only stamp: still comparable, just not provable."""
    if not model_path:
        return {"model_name": "", "model_sha256": "", "embed_dim": embed_dim}
    sha = sha256_file(model_path)
    if not sha:
        print(f"[gallery] cannot hash embedder model {model_path} — model binding "
              f"falls back to filename only (GR-004)", file=sys.stderr)
    return {"model_name": Path(model_path).name, "model_sha256": sha,
            "embed_dim": embed_dim}


def _stamp_empty(s) -> bool:
    return not s or (not s.get("model_name") and not s.get("model_sha256"))


def describe_stamp(s) -> str:
    if _stamp_empty(s):
        return "UNKNOWN"
    name = s.get("model_name") or "<unnamed model>"
    sha = s.get("model_sha256") or ""
    return f"{name} (sha256 {sha[:12]}…)" if sha else f"{name} (sha256 unavailable)"


def require_gallery_stamp_from_env() -> bool:
    """SAE_REQUIRE_GALLERY_STAMP=1 → an unprovable binding is fatal, not a warning."""
    return os.environ.get("SAE_REQUIRE_GALLERY_STAMP", "0") not in ("", "0")


def check_embedder_stamp(built_with: dict | None, loading_with: dict | None,
                         gallery_desc: str = "gallery",
                         embedder_desc: str = "embedder") -> tuple[str, str]:
    """Pure comparison. Returns (verdict, message); verdict is one of
    match / weak_match / unstamped / unknown_embedder / mismatch.

    Same rules as compare_embedder_stamps() in src/gallery/embedder_stamp.cpp."""
    if _stamp_empty(built_with):
        return "unstamped", (
            f"gallery '{gallery_desc}' carries no embedder stamp (GR-004).\n"
            f"  gallery was built with : UNKNOWN — this file predates model binding\n"
            f"  embedder now loaded    : {describe_stamp(loading_with)}  [{embedder_desc}]\n"
            f"  If these are not the same model every similarity from this run is\n"
            f"  meaningless but will look plausible. Rebuild or re-stamp the gallery\n"
            f"  (scripts/stamp_gallery.py), or run with SAE_REQUIRE_GALLERY_STAMP=1 to\n"
            f"  make this a hard error.")

    if _stamp_empty(loading_with):
        return "unknown_embedder", (
            f"cannot identify the embedder being used against gallery "
            f"'{gallery_desc}' (GR-004).\n"
            f"  gallery was built with : {describe_stamp(built_with)}\n"
            f"  embedder now loaded    : UNKNOWN [{embedder_desc}]\n"
            f"  The binding cannot be checked, so it is not being checked.")

    mismatch_tail = (
        "  Cosine similarities between embeddings from different models are\n"
        "  meaningless but look plausible. Rebuild the gallery with the loaded\n"
        "  model, or point the embedder at the model the gallery was built with.")

    if int(built_with.get("embed_dim", 512)) != int(loading_with.get("embed_dim", 512)):
        return "mismatch", (
            f"gallery/embedder MODEL MISMATCH — refusing to run (GR-004).\n"
            f"  gallery was built with : {describe_stamp(built_with)}, "
            f"dim={built_with.get('embed_dim')}  [{gallery_desc}]\n"
            f"  embedder now loaded    : {describe_stamp(loading_with)}, "
            f"dim={loading_with.get('embed_dim')}  [{embedder_desc}]\n"
            f"  Embedding dimensions differ; these are not the same space.")

    a, b = built_with.get("model_sha256", ""), loading_with.get("model_sha256", "")
    if a and b:
        if a == b:
            note = ""
            if built_with.get("model_name") != loading_with.get("model_name"):
                note = (f" (gallery recorded it as '{built_with.get('model_name')}', "
                        f"loaded from '{loading_with.get('model_name')}' — "
                        f"same bytes, renamed file)")
            return "match", f"embedder binding verified: {describe_stamp(built_with)}{note}"
        return "mismatch", (
            f"gallery/embedder MODEL MISMATCH — refusing to run (GR-004).\n"
            f"  gallery was built with : {built_with.get('model_name')}  sha256={a}\n"
            f"                           [{gallery_desc}]\n"
            f"  embedder now loaded    : {loading_with.get('model_name')}  sha256={b}\n"
            f"                           [{embedder_desc}]\n" + mismatch_tail)

    if built_with.get("model_name") and \
            built_with.get("model_name") == loading_with.get("model_name"):
        return "weak_match", (
            f"embedder binding UNPROVEN for gallery '{gallery_desc}' (GR-004).\n"
            f"  gallery was built with : {describe_stamp(built_with)}\n"
            f"  embedder now loaded    : {describe_stamp(loading_with)}  [{embedder_desc}]\n"
            f"  Filenames agree but at least one SHA-256 is unavailable, so an\n"
            f"  in-place re-export under the same name would not be detected.")

    return "mismatch", (
        f"gallery/embedder MODEL MISMATCH — refusing to run (GR-004).\n"
        f"  gallery was built with : {describe_stamp(built_with)}  [{gallery_desc}]\n"
        f"  embedder now loaded    : {describe_stamp(loading_with)}  [{embedder_desc}]\n"
        + mismatch_tail)


def enforce_embedder_stamp(built_with, loading_with, gallery_desc, embedder_desc,
                           require_stamp: bool = False) -> str:
    """Apply check_embedder_stamp: raise EmbedderMismatch when fatal, else warn.

    A mismatch is fatal unconditionally — there is no bypass, because a mismatch is
    a known-wrong state, not an unknown one. The three "cannot prove it" verdicts
    warn loudly and become fatal under require_stamp / SAE_REQUIRE_GALLERY_STAMP."""
    strict = require_stamp or require_gallery_stamp_from_env()
    verdict, msg = check_embedder_stamp(built_with, loading_with,
                                        gallery_desc, embedder_desc)
    if verdict == "mismatch":
        raise EmbedderMismatch(msg)
    if strict and verdict != "match":
        raise EmbedderMismatch(
            msg + "\n  (fatal because SAE_REQUIRE_GALLERY_STAMP is set)")
    if verdict == "match":
        print(f"[gallery] {msg}", file=sys.stderr)
    else:
        print(f"\n[gallery] ***** WARNING (GR-004) *****\n{msg}\n"
              f"[gallery] ****************************\n", file=sys.stderr)
    return verdict


def read_gallery_stamp(path) -> dict | None:
    """The embedder stamp recorded in a gallery file, or None if unstamped.

    Handles both the HDF5 /embedder group and the legacy JSON "embedder" object."""
    path = Path(path)
    if path.suffix in (".h5", ".hdf5"):
        with h5py.File(path, "r") as f:
            if "embedder" not in f:
                return None
            a = f["embedder"].attrs
            return {"model_name": _as_str(a.get("model_name", "")),
                    "model_sha256": _as_str(a.get("model_sha256", "")),
                    "embed_dim": int(a.get("embed_dim", 512))}
    data = json.loads(path.read_text())
    return data.get("embedder") or None


def verify_gallery_stamp(gallery_path, model_path=None, *, stamp=None,
                         embedder_desc: str | None = None,
                         require_stamp: bool = False) -> str:
    """Load a gallery's stamp and check it against a model file (or an explicit
    stamp, e.g. one read off an embedding dump). Raises EmbedderMismatch."""
    loading = stamp if stamp is not None else embedder_stamp(model_path)
    return enforce_embedder_stamp(read_gallery_stamp(gallery_path), loading,
                                  str(gallery_path),
                                  embedder_desc or str(model_path or "unknown"),
                                  require_stamp)


def _as_str(v) -> str:
    return v.decode() if isinstance(v, bytes) else ("" if v is None else str(v))