train/dist_align_gate.py

"""dist_align_gate.py — teacher-bank alignment verification + model
loading for the distillation bed.

The bank stores UNNORMALIZED pooled image features; every consumer
L2-normalizes at load. Feature parity requires exact preprocessing:
plain-GELU CLIP forward (NOT QuickGELU — the open_clip default for some
checkpoints differs at the 0.975-cosine level), torchvision
Resize(224, bicubic shortest-edge) + CenterCrop(224) + CLIP mean/std,
CLS -> projection readout, fp32 compute. SigLIP uses its own processor
(384x384 warp). Two hub checkpoints ship open_clip-format weights only
(LAION B/16, DataComp B/32); this module converts them deterministically
into transformers CLIPModel format, and reproduction of the stored bank
features (cos >= 0.9999 on reference rows) is the proof of conversion.
"""
import argparse
import glob
import json
import os
import sys
import zlib

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
DEV = "cuda" if torch.cuda.is_available() else "cpu"

DATA = os.path.join(os.environ.get("DIST_DATA_ROOT", "./data"),
                    "dist")
BANK = os.path.join(DATA, "bank")
VAL_DIR = os.path.join(DATA, "coco", "val2017")
CAPS_JSON = os.path.join(DATA, "coco", "annotations", "captions_val2017.json")

GATE_CONFIGS = ("clip_b16_laion2b", "siglip_b16_384", "clip_b32_openai")
CFG_MODEL = {
    "clip_b16_laion2b": "laion/CLIP-ViT-B-16-laion2B-s34B-b88K",
    "siglip_b16_384": "google/siglip-base-patch16-384",
    "clip_b32_openai": "openai/clip-vit-base-patch32",
    "clip_b32_laion2b": "laion/CLIP-ViT-B-32-laion2B-s34B-b79K",
    "clip_b32_datacomp": "laion/CLIP-ViT-B-32-DataComp.XL-s13B-b90K",
    "clip_b16_openai": "openai/clip-vit-base-patch16",
    "dinov2_b14": "facebook/dinov2-base",
}
PRIMARY = "clip_b16_laion2b"
LAION_HF_DIR = os.path.join(DATA, "hf_clip_b16_laion2b")
DATACOMP_HF_DIR = os.path.join(DATA, "hf_clip_b32_datacomp")


def seed_for(name):
    return zlib.crc32(name.encode()) & 0x7FFFFFFF


def _convert_open_clip(model_id, out_dir, patch):
    """open_clip-only hub checkpoints (laion b16, datacomp b32 — no
    transformers config.json, verified) -> deterministic key-map into an HF
    CLIPModel, saved once to `out_dir`. The bank-parity gate is the
    end-to-end verifier (a wrong map cannot produce cos ~1 vs stored bank
    features). ViT-B geometry is shared; only patch_size differs."""
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file
    from transformers import AutoTokenizer, CLIPConfig, CLIPImageProcessor, CLIPModel
    w = load_file(hf_hub_download(model_id, "open_clip_model.safetensors"))
    cfg = CLIPConfig(
        projection_dim=512,
        # laion-family open_clip ViT-B = PLAIN gelu (quick_gelu is the
        # OpenAI-checkpoint quirk and HF CLIPConfig's default — must override)
        vision_config=dict(hidden_size=768, intermediate_size=3072,
                           num_hidden_layers=12, num_attention_heads=12,
                           image_size=224, patch_size=patch,
                           projection_dim=512, hidden_act="gelu"),
        text_config=dict(hidden_size=512, intermediate_size=2048,
                         num_hidden_layers=12, num_attention_heads=8,
                         max_position_embeddings=77, vocab_size=49408,
                         projection_dim=512, hidden_act="gelu"))
    model = CLIPModel(cfg)
    sd = {}

    def blocks(src, dst, d):
        for i in range(12):
            p, q = f"{src}.resblocks.{i}.", f"{dst}.layers.{i}."
            for s, t in (("ln_1", "layer_norm1"), ("ln_2", "layer_norm2"),
                         ("attn.out_proj", "self_attn.out_proj"),
                         ("mlp.c_fc", "mlp.fc1"), ("mlp.c_proj", "mlp.fc2")):
                sd[q + t + ".weight"] = w[p + s + ".weight"]
                sd[q + t + ".bias"] = w[p + s + ".bias"]
            for j, t in enumerate(("q_proj", "k_proj", "v_proj")):
                sd[q + f"self_attn.{t}.weight"] = \
                    w[p + "attn.in_proj_weight"][j * d:(j + 1) * d]
                sd[q + f"self_attn.{t}.bias"] = \
                    w[p + "attn.in_proj_bias"][j * d:(j + 1) * d]

    blocks("visual.transformer", "vision_model.encoder", 768)
    blocks("transformer", "text_model.encoder", 512)
    sd["vision_model.embeddings.class_embedding"] = w["visual.class_embedding"]
    sd["vision_model.embeddings.position_embedding.weight"] = \
        w["visual.positional_embedding"]
    sd["vision_model.embeddings.patch_embedding.weight"] = w["visual.conv1.weight"]
    pre = [k[:-7] for k in model.state_dict()
           if "pre_lay" in k and k.endswith(".weight")][0]   # pre_layrnorm typo-BC
    sd[pre + ".weight"] = w["visual.ln_pre.weight"]
    sd[pre + ".bias"] = w["visual.ln_pre.bias"]
    sd["vision_model.post_layernorm.weight"] = w["visual.ln_post.weight"]
    sd["vision_model.post_layernorm.bias"] = w["visual.ln_post.bias"]
    sd["visual_projection.weight"] = w["visual.proj"].t().contiguous()
    sd["text_model.embeddings.token_embedding.weight"] = w["token_embedding.weight"]
    sd["text_model.embeddings.position_embedding.weight"] = w["positional_embedding"]
    sd["text_model.final_layer_norm.weight"] = w["ln_final.weight"]
    sd["text_model.final_layer_norm.bias"] = w["ln_final.bias"]
    sd["text_projection.weight"] = w["text_projection"].t().contiguous()
    sd["logit_scale"] = w["logit_scale"].reshape(())
    missing, unexpected = model.load_state_dict(sd, strict=False)
    missing = [m for m in missing if "position_ids" not in m]
    assert not missing and not unexpected, (missing, unexpected)
    os.makedirs(out_dir, exist_ok=True)
    model.save_pretrained(out_dir)
    AutoTokenizer.from_pretrained(CFG_MODEL[PRIMARY]).save_pretrained(out_dir)
    CLIPImageProcessor().save_pretrained(out_dir)   # defaults == the repo's
    # open_clip preprocess_cfg: 224 bicubic shortest-edge + center-crop,
    # OpenAI CLIP mean/std (verified against open_clip_config.json)
    print(f"[GATE] converted open_clip -> HF CLIPModel at {out_dir}")


def laion_local_dir():
    """Path to the transformers-format LAION B/16 (converted once, cached).
    dist_bed.py imports this for its text/vision teacher heads."""
    if not os.path.isfile(os.path.join(LAION_HF_DIR, "model.safetensors")):
        _convert_open_clip(CFG_MODEL[PRIMARY], LAION_HF_DIR, 16)
    return LAION_HF_DIR


def datacomp_local_dir():
    """Transformers-format DataComp-XL B/32 (open_clip-only upstream)."""
    if not os.path.isfile(os.path.join(DATACOMP_HF_DIR, "model.safetensors")):
        _convert_open_clip(CFG_MODEL["clip_b32_datacomp"],
                           DATACOMP_HF_DIR, 32)
    return DATACOMP_HF_DIR


def load_bank(config, split="val"):
    """-> (ids int64 np, feats float32 torch UNNORMALIZED) in parquet row order."""
    import pyarrow as pa
    import pyarrow.parquet as pq
    files = sorted(glob.glob(os.path.join(BANK, config, f"{split}-*.parquet")))
    if not files:
        print(f"[GATE] MISSING bank parquet: {config}/{split} under {BANK}")
        sys.exit(2)
    t = pa.concat_tables([pq.read_table(f, columns=["image_id", "features"])
                          for f in files])
    ids = t.column("image_id").to_numpy()
    fl = t.column("features").combine_chunks()
    D = fl.type.list_size
    feats = fl.flatten().to_numpy(zero_copy_only=False).reshape(len(ids), D)
    return ids, torch.from_numpy(np.array(feats, dtype=np.float32, copy=True))


def load_vision(config):
    """(model, image_processor) for a bank config, fp32 eval on DEV. The LAION
    primary loads from the local converted dir (open_clip-only upstream).
    Image processor only — SigLIP's slow tokenizer needs sentencepiece (absent
    in this venv); the image side never touches a tokenizer.

    THE BANK ACT FINDING (measured 2026-07-26, this gate's discovery run):
    the bank's CLIP towers were embedded with PLAIN GELU — for the OpenAI
    checkpoints that is the open_clip create_model('ViT-B-32', 'openai')
    silent QuickGELU mismatch (plain-gelu local re-embed: cos mean 1.00000
    min 0.99999 vs quick_gelu 0.97515; geometry/precision/pooling/crop all
    falsified first). Preprocessing = the default 224 shortest-edge bicubic
    + center-crop + CLIP mean/std; siglip = its default 384 warp; storage is
    fp16-precision in fp32. The bed MUST match the bank, so every CLIP tower
    here runs hidden_act='gelu' (correct anyway for laion2b checkpoints)."""
    from transformers import AutoConfig, AutoImageProcessor, AutoModel
    if config == PRIMARY:
        name = laion_local_dir()
    elif config == "clip_b32_datacomp":
        name = datacomp_local_dir()
    else:
        name = CFG_MODEL[config]
    kw, proc = {}, None
    if config.startswith("clip_"):
        c = AutoConfig.from_pretrained(name)
        c.vision_config.hidden_act = "gelu"
        c.text_config.hidden_act = "gelu"
        kw["config"] = c
        # bank CLIP geometry is torchvision-exact (measured 1.00000/0.99999
        # on b32 vs 0.99919 through the HF processor — crop-rounding differs)
        proc = clip_torchvision_pipe(224)
    else:
        proc = AutoImageProcessor.from_pretrained(name)
    model = _from_pretrained_fp32(name, **kw).to(DEV).eval()
    return model, proc


def _from_pretrained_fp32(name, **kw):
    """transformers 5.x takes dtype=, 4.x takes torch_dtype= — accept both
    (the pod runs 4.49 pinned under torch 2.4.1; local runs 5.x)."""
    from transformers import AutoModel
    try:
        return AutoModel.from_pretrained(name, dtype=torch.float32, **kw)
    except TypeError:
        return AutoModel.from_pretrained(name, torch_dtype=torch.float32,
                                         **kw)


CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
CLIP_STD = (0.26862954, 0.26130258, 0.27577711)


class clip_torchvision_pipe:
    """Callable matching the AutoImageProcessor interface for the ONE call
    shape this gate uses: proc(images=[PIL...], return_tensors='pt')."""
    def __init__(self, size):
        from torchvision import transforms as T
        from torchvision.transforms import InterpolationMode
        self.tf = T.Compose([
            T.Resize(size, interpolation=InterpolationMode.BICUBIC),
            T.CenterCrop(size), T.ToTensor(),
            T.Normalize(CLIP_MEAN, CLIP_STD)])

    def __call__(self, images, return_tensors="pt", **kw):
        return {"pixel_values": torch.stack([self.tf(im) for im in images])}


def load_siglip_text_tokenizer():
    """SigLIP fast tokenizer from the repo's own tokenizer.json (no
    sentencepiece dependency). Exported for dist_bed.py's siglip text head."""
    from huggingface_hub import hf_hub_download
    from transformers import PreTrainedTokenizerFast
    tj = hf_hub_download(CFG_MODEL["siglip_b16_384"], "tokenizer.json")
    return PreTrainedTokenizerFast(tokenizer_file=tj, pad_token="</s>",
                                   eos_token="</s>", unk_token="<unk>",
                                   model_max_length=64)


def _feat(out):
    """transformers 5.x returns the full output object from get_*_features
    (projected features live in .pooler_output); 4.x returned the tensor."""
    return out if torch.is_tensor(out) else out.pooler_output


@torch.no_grad()
def embed_images(model, proc, pils, bs=16):
    outs = []
    for i in range(0, len(pils), bs):
        px = proc(images=pils[i:i + bs], return_tensors="pt")["pixel_values"]
        outs.append(_feat(model.get_image_features(
            pixel_values=px.to(DEV, torch.float32))).float().cpu())
    return torch.cat(outs)


@torch.no_grad()
def embed_texts(model, tok, texts, bs=64, siglip=False):
    outs = []
    for i in range(0, len(texts), bs):
        kw = (dict(padding="max_length", truncation=True, max_length=64)
              if siglip else dict(padding=True, truncation=True))
        tk = tok(texts[i:i + bs], return_tensors="pt", **kw)
        tk = {k: v.to(DEV) for k, v in tk.items()
              if k in ("input_ids", "attention_mask")}
        outs.append(_feat(model.get_text_features(**tk)).float().cpu())
    return torch.cat(outs)


def first_captions():
    if not os.path.isfile(CAPS_JSON):
        print(f"[GATE] MISSING captions json: {CAPS_JSON}")
        sys.exit(2)
    with open(CAPS_JSON, "r", encoding="utf-8") as f:
        anns = json.load(f)["annotations"]
    cap = {}
    for a in anns:                       # first caption in annotation order
        cap.setdefault(int(a["image_id"]), a["caption"].strip())
    return cap


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=64)
    ap.add_argument("--n-text", type=int, default=256)
    ap.add_argument("--margin", type=float, default=0.05)
    a, _ = ap.parse_known_args()

    if DEV == "cuda":
        torch.cuda.set_per_process_memory_fraction(0.73)
    if not os.path.isdir(VAL_DIR) or not glob.glob(os.path.join(VAL_DIR, "*.jpg")):
        print(f"[GATE] val2017 images absent/empty at {VAL_DIR} — still "
              f"downloading? Re-run once populated.")
        sys.exit(2)

    rows, any_fail = [], False
    for cfg in GATE_CONFIGS:
        ids, feats = load_bank(cfg)
        g = torch.Generator().manual_seed(seed_for(f"dist-align:{cfg}"))
        sel = torch.randperm(len(ids), generator=g)[:a.n].numpy()
        pils = []
        for i in sel:
            p = os.path.join(VAL_DIR, f"{int(ids[i]):012d}.jpg")
            if not os.path.isfile(p):
                print(f"[GATE] missing jpg {p} — val2017 incomplete?")
                sys.exit(2)
            pils.append(Image.open(p).convert("RGB"))
        model, proc = load_vision(cfg)
        loc = F.normalize(embed_images(model, proc, pils), dim=-1)
        bank = F.normalize(feats[sel], dim=-1)
        cos = (loc * bank).sum(-1)
        mean_c, min_c = float(cos.mean()), float(cos.min())
        verdict = ("PASS" if mean_c >= 0.999 else
                   "WARN" if mean_c >= 0.99 else "FAIL")
        any_fail |= verdict == "FAIL"
        rows.append((cfg, feats.shape[1], len(sel), mean_c, min_c, verdict))
        del model, proc
        if DEV == "cuda":
            torch.cuda.empty_cache()

    # ---- LAION text-head sanity ------------------------------------------
    from transformers import AutoModel, AutoTokenizer
    ldir = laion_local_dir()
    tok = AutoTokenizer.from_pretrained(ldir)
    model = _from_pretrained_fp32(ldir).to(DEV).eval()
    probes = ["a photo of a cat", "a photo of a dog", "a photo of a car"]
    praw = embed_texts(model, tok, probes)
    raw_norms = praw.norm(dim=-1)
    pnorm = F.normalize(praw, dim=-1)
    norm_dev = float((pnorm.norm(dim=-1) - 1.0).abs().max())
    ids, feats = load_bank(PRIMARY)
    cap = first_captions()
    g = torch.Generator().manual_seed(seed_for("dist-align-text"))
    sel = torch.randperm(len(ids), generator=g)[:a.n_text].numpy()
    texts = [cap[int(ids[i])] for i in sel]
    temb = F.normalize(embed_texts(model, tok, texts), dim=-1)
    img = F.normalize(feats[sel], dim=-1)
    matched = float((temb * img).sum(-1).mean())
    shuffled = float((temb.roll(1, 0) * img).sum(-1).mean())
    margin = matched - shuffled
    tverdict = ("PASS" if (margin >= a.margin and norm_dev < 1e-3
                           and bool(torch.isfinite(praw).all())) else "FAIL")
    any_fail |= tverdict == "FAIL"

    print("=== DIST ALIGNMENT GATE (local re-embed vs stored bank, val2017) ===")
    print(f"{'config':<18} {'dim':>4} {'n':>4} {'mean_cos':>9} {'min_cos':>9} verdict")
    for cfg, D, n, mc, mn, v in rows:
        print(f"{cfg:<18} {D:>4} {n:>4} {mc:>9.5f} {mn:>9.5f} {v}")
    print(f"text-head sanity (laion b16): probe raw norms "
          f"[{float(raw_norms.min()):.3f}..{float(raw_norms.max()):.3f}] "
          f"normalized max|1-n| {norm_dev:.1e}; "
          f"caption cos matched {matched:.4f} vs shuffled {shuffled:.4f} "
          f"margin {margin:.4f} (bar {a.margin})  {tverdict}")
    overall = "FAIL" if any_fail else "PASS"
    print(f"OVERALL: {overall}")
    sys.exit(1 if any_fail else 0)


if __name__ == "__main__":
    main()