Blame for 45b_text_feats_pathvlm.py

"""Frozen TEXT-tower embeddings of two pathology vision-language models (same output format as 45_text_feats.py).

    python 45b_text_feats_pathvlm.py <name> [<name> ...]     # name in {keep, quiltnet, quiltnet-gelu}
    python 45b_text_feats_pathvlm.py verify [keep|quiltnet]    # checks against the official code (see below)

Text sets / row order / output exactly as 45_text_feats.py: cache/<set>_text_<name>.npy, (n, 1, D) float16, raw
(pre-normalisation) embedding, for alltxt, orion_gentext, tcga_notext, dlpfc_gentext (column `caption`).
Inference in fp32 (no autocast; the models are small) and cast to float16 at save time.
Runtime (shared A5000, 1 loaded CPU): keep 85 s for all 4 sets, quiltnet + quiltnet-gelu 35 s.
Leakage check (papers + repos, 2026-10-24): neither model was pretrained on PathGen-1.6M or other TCGA captions.
  KEEP: OpenPath (Twitter + LAION) + Quilt-1M, regrouped into 143k semantic groups (Loie/KEEP_dataset); text
  encoder PubMedBERT-arch., initialised from a knowledge encoder trained on their disease KG (UMLS-based).
  QuiltNet: OpenAI CLIP fine-tuned on Quilt-1M = YouTube 802k + PubMed OA 59k + Twitter/OpenPath list ~134k
  + LAION 23k (paper June 2023, a year before PathGen-1.6M, July 2024). TCGA appears only as evaluation data.

keep      KEEP (Zhou et al., Cancer Cell 2026; github MAGIC-AI4Med/KEEP), HF Astaxanthin/KEEP
          revision 28a25d95cc6ba27a7e6fab3f144e13dbafd8b21e, licence MIT. D = 768.
          Text tower = PubMedBERT-architecture BertModel (12 layers, 768, vocab 30522) incl. its BERT pooler
          (tanh(Linear(h_CLS))); KEEP has no separate text projection - its official modeling_keep.py
          encode_text() is normalize(BertModel(**tokens).pooler_output). We return pooler_output before the
          L2-normalisation. Tokenisation as in the README / repo zero-shot example: AutoTokenizer of the repo
          (BertTokenizerFast, repo vocab.txt, lower-case), max_length=256, padding='max_length', truncation=True (input_ids, token_type_ids,
          attention_mask all passed). Weights: text.* of model.safetensors, strict load. The official
          modeling_keep.py (copied to ext/keep/, read in full: model definition only, but it imports timm - not in
          scileo - and monkey-patches timm's LayerScale) is therefore NOT imported for extraction; the text tower is
          rebuilt here with transformers' BertModel.
quiltnet  QuiltNet-B-32 (Ikezogwo et al., NeurIPS 2023 D&B), HF wisdomik/QuiltNet-B-32
          revision 8ce77289ce35a90b2f1db1137dfa4bc2df175e33, licence MIT. D = 512.
          open_clip CLIP text tower (12 layers, width 512, 8 heads, causal mask, 77 ctx) re-implemented with
          transformers' CLIPTextModelWithProjection and the open_clip state dict (open_clip_pytorch_model.bin)
          remapped key-by-key (in_proj split into q/k/v, text_projection transposed), strict load.
          Activation = QuickGELU. QuiltNet was fine-tuned from OpenAI CLIP ViT-B/32 with OpenCLIP (paper sec. 4 / app.
          C.2), and OpenAI CLIP uses QuickGELU, but the repo's open_clip_config.json has no quick_gelu flag (and its HF
          config.json says hidden_act 'gelu'), so a plain open_clip 'hf-hub:' load silently runs GELU. On image<->caption
          retrieval with the real open_clip model (both towers), QuickGELU beats GELU on every metric: TCGA 256 i2t R@1
          .082 vs .051, t2i .094 vs .066; ORION 512 i2t .029 vs .012, t2i .031 vs .021. So the weights were trained
          with QuickGELU and the released config is wrong; we use QuickGELU (= open_clip
          create_model('hf-hub:wisdomik/QuiltNet-B-32', force_quick_gelu=True)). The literal released-config (GELU)
          variant is available as name 'quiltnet-gelu' (mean cosine to the QuickGELU embedding only ~0.78).
          Tokenisation = open_clip's default SimpleTokenizer, reproduced as: whitespace_clean(basic_clean(t)).lower()
          (ftfy.fix_text 6.3.1 from ext/quiltnet/pylib (pure python, not installed), html.unescape x2), OpenAI CLIP BPE
          (CLIPTokenizerFast of openai/clip-vit-base-patch32), [SOT] + first 75 BPE + [EOT], zero padding to 77, no
          padding mask. Pooling = hidden state at the EOT token (argmax of ids) after ln_final, @ text_projection
          = open_clip encode_text(normalize=False).

Verification (`verify`, run 2026-10-24; references loaded from unpacked wheels in ext/*/pylib, never pip-installed):
  keep     official KEEPModel (ext/keep/modeling_keep.py, timm 1.0.15 from ext/keep/pylib, from_pretrained: all keys
           matched): its encode_text == normalize(ours) BIT-EXACT (max|diff| 0) on 67 texts (64 real captions, an empty
           and an over-length one); README AutoTokenizer == slow BertTokenizer inputs on 10405/10405 captions; README
           example.tif zero-shot: breast invasive carcinoma 0.79 / lung adeno 0.21 / normal 0.001; 256-tile TCGA
           image<->synthetic caption retrieval with the official image tower i2t R@1 0.094, t2i 0.063 (chance 0.004).
  quiltnet open_clip 2.32.0 create_model('hf-hub:...', force_quick_gelu=True).encode_text vs ours (fp32, CPU):
           max|diff| ~1e-5 (|x| up to ~8), min cosine 0.9999998 (float-op order only); same for the GELU variant vs the
           default open_clip load; open_clip token ids == ours on 10405/10405 captions; remapped weights == the repo's
           own HF-format pytorch_model.bin text weights (197/197 tensors, max|diff| 0); repo ADI-TCGA example tile
           zero-shot over 9 CRC tissue classes -> adipose (top-1); retrieval numbers above.
"""
import html
import json
import os
import re
import sys
import time

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download

from common import CACHE, ROOT

dev = "cuda"
SETS = {"alltxt": "alltxt_meta.csv", "orion_gentext": "orion_gentext_meta.csv", "tcga_notext": "tcga_notext_meta.csv",
        "dlpfc_gentext": "dlpfc_gentext_meta.csv"}
KEEP_REPO, KEEP_REV = "Astaxanthin/KEEP", "28a25d95cc6ba27a7e6fab3f144e13dbafd8b21e"
QN_REPO, QN_REV = "wisdomik/QuiltNet-B-32", "8ce77289ce35a90b2f1db1137dfa4bc2df175e33"
CLIP_TOK = "openai/clip-vit-base-patch32"
EXT = os.path.join(ROOT, "ext")


class KEEPText(nn.Module):
    """KEEP text tower: BertModel (+pooler) -> pooler_output (un-normalised)."""

    def __init__(self):
        super().__init__()
        from safetensors.torch import load_file
        from transformers import AutoTokenizer, BertConfig, BertModel
        cfg = json.load(open(hf_hub_download(KEEP_REPO, "config.json", revision=KEEP_REV)))
        self.text = BertModel(BertConfig(**cfg["text_config"]), add_pooling_layer=True)
        sd = load_file(hf_hub_download(KEEP_REPO, "model.safetensors", revision=KEEP_REV))
        sd = {k[5:]: v for k, v in sd.items() if k.startswith("text.")}
        self.text.load_state_dict(sd, strict=True)
        self.tok = AutoTokenizer.from_pretrained(KEEP_REPO, revision=KEEP_REV)  # as README (BertTokenizer(Fast), lower)

    def tokenize(self, texts):
        return self.tok(texts, max_length=256, padding="max_length", truncation=True, return_tensors="pt")

    def forward(self, texts):
        x = self.tokenize(texts).to(next(self.parameters()).device)
        return self.text(**x).pooler_output


def _oc_clean(t):  # = open_clip tokenizer 'lower' clean: whitespace_clean(basic_clean(t)).lower()
    import ftfy  # ftfy 6.3.1 (pure python, unpacked wheel in ext/quiltnet/pylib; e.g. uncurls quotes as open_clip does)
    t = html.unescape(html.unescape(ftfy.fix_text(t))).strip()
    return re.sub(r"\s+", " ", t).strip().lower()


class QuiltText(nn.Module):
    """open_clip CLIP text tower of QuiltNet-B-32 on transformers' CLIPTextModelWithProjection."""
    CTX = 77

    def __init__(self, act="quick_gelu"):
        super().__init__()
        from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizerFast
        sys.path.append(os.path.join(EXT, "quiltnet", "pylib"))  # for ftfy only (open_clip's text cleaner)
        oc = json.load(open(hf_hub_download(QN_REPO, "open_clip_config.json", revision=QN_REV)))["model_cfg"]
        tc = oc["text_cfg"]
        cfg = CLIPTextConfig(vocab_size=tc["vocab_size"], hidden_size=tc["width"], intermediate_size=4 * tc["width"],
                             num_hidden_layers=tc["layers"], num_attention_heads=tc["heads"],
                             max_position_embeddings=tc["context_length"], projection_dim=oc["embed_dim"],
                             hidden_act=act, layer_norm_eps=1e-5, eos_token_id=2)  # eos_token_id=2 -> argmax pooling
        self.m = CLIPTextModelWithProjection(cfg)
        sd = torch.load(hf_hub_download(QN_REPO, "open_clip_pytorch_model.bin", revision=QN_REV), map_location="cpu",
                        weights_only=True)
        self.m.load_state_dict(self.map_open_clip(sd, tc["layers"]), strict=True)
        self.tok = CLIPTokenizerFast.from_pretrained(CLIP_TOK)
        self.sot, self.eot = self.tok.bos_token_id, self.tok.eos_token_id  # 49406, 49407

    @staticmethod
    def map_open_clip(sd, L):
        t = "text_model."
        o = {t + "embeddings.token_embedding.weight": sd["token_embedding.weight"],
             t + "embeddings.position_embedding.weight": sd["positional_embedding"],
             t + "final_layer_norm.weight": sd["ln_final.weight"], t + "final_layer_norm.bias": sd["ln_final.bias"],
             "text_projection.weight": sd["text_projection"].t().contiguous()}
        for i in range(L):
            a, b = f"transformer.resblocks.{i}.", f"{t}encoder.layers.{i}."
            for w in ("weight", "bias"):
                q, k, v = sd[a + "attn.in_proj_" + w].chunk(3, 0)
                o.update({b + "self_attn.q_proj." + w: q, b + "self_attn.k_proj." + w: k, b + "self_attn.v_proj." + w: v,
                          b + "self_attn.out_proj." + w: sd[a + "attn.out_proj." + w],
                          b + "layer_norm1." + w: sd[a + "ln_1." + w], b + "layer_norm2." + w: sd[a + "ln_2." + w],
                          b + "mlp.fc1." + w: sd[a + "mlp.c_fc." + w], b + "mlp.fc2." + w: sd[a + "mlp.c_proj." + w]})
        return o

    def tokenize(self, texts):
        out = torch.zeros(len(texts), self.CTX, dtype=torch.long)
        for i, ids in enumerate(self.tok([_oc_clean(t) for t in texts], add_special_tokens=False).input_ids):
            ids = [self.sot] + ids[:self.CTX - 2] + [self.eot]
            out[i, :len(ids)] = torch.tensor(ids)
        return out

    def forward(self, texts):
        ids = self.tokenize(texts).to(next(self.parameters()).device)
        return self.m(input_ids=ids).text_embeds  # EOT pooling (argmax id), ln_final, @ text_projection


def encoder(name):
    m = {"keep": KEEPText, "quiltnet": lambda: QuiltText("quick_gelu"),
         "quiltnet-gelu": lambda: QuiltText("gelu")}[name]().to(dev).eval()
    return m, 64


def load_texts(meta):
    return pd.read_csv(os.path.join(CACHE, meta)).caption.fillna("").astype(str).tolist()


def extract(names):
    for name in names:
        t0 = time.time()
        f, bs = encoder(name)
        for st, meta in SETS.items():
            out_f = os.path.join(CACHE, f"{st}_text_{name}.npy")
            if os.path.exists(out_f):
                continue
            txt = load_texts(meta)
            order = np.argsort([len(t) for t in txt])  # length-sorted batches, restored below
            E = None
            with torch.no_grad():
                for i in range(0, len(txt), bs):
                    idx = order[i:i + bs]
                    z = f([txt[j] for j in idx]).float().cpu().numpy()
                    if E is None:
                        E = np.zeros((len(txt), z.shape[1]), np.float32)
                    E[idx] = z
            assert np.isfinite(E).all()
            np.save(out_f, E[:, None].astype(np.float16))
            print(f"[{name}] {st}: {E.shape}  {time.time() - t0:.0f}s", flush=True)
        del f
        torch.cuda.empty_cache()


# ----------------------------------------------------------------------------------------------- verification
def _retrieval(img_feats, txt_feats):
    a = torch.nn.functional.normalize(img_feats, dim=-1)
    b = torch.nn.functional.normalize(txt_feats, dim=-1)
    S = a @ b.T
    r = torch.arange(len(S))
    return {"i2t_R@1": (S.argmax(1) == r).float().mean().item(), "t2i_R@1": (S.argmax(0) == r).float().mean().item(),
            "i2t_R@5": (S.topk(5, 1).indices == r[:, None]).any(1).float().mean().item()}


def verify(which=("keep", "quiltnet")):
    from PIL import Image, PngImagePlugin
    PngImagePlugin.MAX_TEXT_CHUNK = 1 << 30  # some PathGen PNGs carry huge iCCP chunks
    torch.manual_seed(0)
    texts = load_texts(SETS["alltxt"])
    sample = [texts[i] for i in np.random.default_rng(0).choice(len(texts), 64, replace=False)]
    sample += ["an H&E image of breast invasive carcinoma.", "", "a " + "very " * 120 + "long text"]
    tn = pd.read_csv(os.path.join(CACHE, SETS["tcga_notext"]))
    tn = tn[tn.caption.notna()].reset_index(drop=True)
    imgs = [Image.open(p).convert("RGB") for p in tn.image_path]
    print(f"retrieval set: {len(imgs)} TCGA tiles with synthetic captions (tcga_notext)")
    if "keep" in which:
        verify_keep(sample, tn, imgs)
    if "quiltnet" in which:
        verify_quiltnet(sample, tn, imgs)


def verify_keep(sample, tn, imgs):
    """KEEP: official modeling_keep.py from the local copy, timm from the unpacked wheel."""
    from PIL import Image
    sys.path[:0] = [os.path.join(EXT, "keep", "pylib"), os.path.join(EXT, "keep")]
    from torchvision import transforms
    from transformers import AutoTokenizer
    import modeling_keep  # ext/keep/modeling_keep.py (read in full before use)
    ref, info = modeling_keep.KEEPModel.from_pretrained(KEEP_REPO, revision=KEEP_REV, output_loading_info=True)
    ref = ref.eval()
    print(f"[keep] official from_pretrained loading info: { {k: v for k, v in info.items() if v} or 'all keys matched'}")
    rtok = AutoTokenizer.from_pretrained(os.path.join(EXT, "keep"), use_fast=False)  # slow BertTokenizer as reference
    ours = KEEPText().eval()
    with torch.no_grad():
        a = ref.encode_text(rtok(sample, max_length=256, padding="max_length", truncation=True, return_tensors="pt"))
        b = torch.nn.functional.normalize(ours(sample), dim=-1)
        print(f"[keep] official encode_text vs normalize(ours): max|diff| = {(a - b).abs().max().item():.3e}  "
              f"bit-exact: {torch.equal(a, b)}")
        n_same = n_tok = 0
        for meta in SETS.values():
            t = load_texts(meta)
            for i in range(0, len(t), 256):
                kw = dict(max_length=256, padding="max_length", truncation=True, return_tensors="pt")
                x, y = rtok(t[i:i + 256], **kw), ours.tokenize(t[i:i + 256])
                same = torch.stack([(x[k] == y[k]).all(1) for k in ("input_ids", "token_type_ids", "attention_mask")]).all(0)
                n_tok, n_same = n_tok + len(same), n_same + int(same.sum())
        print(f"[keep] slow BertTokenizer vs our (README AutoTokenizer) inputs identical on {n_same}/{n_tok} captions")
        tf = transforms.Compose([transforms.Resize(224, interpolation=transforms.InterpolationMode.BICUBIC),
                                 transforms.CenterCrop(224), transforms.ToTensor(),
                                 transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
        ex = ["an H&E image of breast invasive carcinoma.", "an H&E image of normal tissue.",
              "an H&E image of lung adenocarcinoma."]
        img = tf(Image.open(os.path.join(EXT, "keep", "example.tif")).convert("RGB"))[None]
        p = (ref.logit_scale.exp() * ref.encode_image(img) @ torch.nn.functional.normalize(ours(ex), dim=-1).T).softmax(-1)
        print(f"[keep] README example.tif zero-shot probs {dict(zip(ex, p[0].round(decimals=3).tolist()))}")
        ref = ref.to(dev)
        I = torch.cat([ref.encode_image(torch.stack([tf(x) for x in imgs[i:i + 32]]).to(dev)).cpu()
                       for i in range(0, len(imgs), 32)])
        T = torch.cat([ours(list(tn.caption[i:i + 64])) for i in range(0, len(tn), 64)])
        print(f"[keep] image<->caption retrieval (chance R@1 {1 / len(imgs):.3f}): {_retrieval(I, T)}")
    del ref, ours
    torch.cuda.empty_cache()


def verify_quiltnet(sample, tn, imgs):
    """QuiltNet: real open_clip 2.32.0 (+ftfy) from the unpacked wheels."""
    from PIL import Image
    sys.path[:0] = [os.path.join(EXT, "quiltnet", "pylib")]
    import open_clip
    from transformers import CLIPModel
    otok = open_clip.get_tokenizer("hf-hub:" + QN_REPO)
    ours = QuiltText().eval()
    n_tok = n_same = 0
    for meta in SETS.values():
        t = load_texts(meta)
        for i in range(0, len(t), 256):
            same = (otok(t[i:i + 256]) == ours.tokenize(t[i:i + 256])).all(1)
            n_tok, n_same = n_tok + len(same), n_same + int(same.sum())
    print(f"[quiltnet] open_clip vs our token ids identical on {n_same}/{n_tok} captions of the 4 sets")
    hf = CLIPModel.from_pretrained(QN_REPO, revision=QN_REV)
    hsd, osd = hf.state_dict(), ours.m.state_dict()
    dw = max((hsd[k] - osd[k]).abs().max().item() for k in osd if k in hsd)
    print(f"[quiltnet] repo HF-format pytorch_model.bin: {sum(k in hsd for k in osd)}/{len(osd)} text tensors present, "
          f"max|w diff| vs remapped open_clip weights = {dw:.3e}; its config hidden_act = {hf.config.text_config.hidden_act}")
    del hf
    cls = ["adipose", "background", "debris", "lymphocytes", "mucus", "smooth muscle", "normal colon mucosa",
           "cancer-associated stroma", "colorectal adenocarcinoma epithelium"]
    prompts = [f"a histopathology image of {c}" for c in cls]
    for act, qg in (("quick_gelu", True), ("gelu", False)):
        oc, _, pre = open_clip.create_model_and_transforms("hf-hub:" + QN_REPO, force_quick_gelu=qg)
        oc = oc.eval()
        mine = ours if qg else QuiltText("gelu").eval()
        with torch.no_grad():
            a, b = oc.encode_text(otok(sample)), mine(sample)
            print(f"[quiltnet/{act}] open_clip encode_text vs ours: max|diff| = {(a - b).abs().max().item():.3e} "
                  f"(max|a| {a.abs().max().item():.2f}), min cos = "
                  f"{torch.nn.functional.cosine_similarity(a, b).min().item():.7f}")
            x = pre(Image.open(os.path.join(EXT, "quiltnet", "ADI-TCGA-CFQTRNMG.tif")).convert("RGB"))[None]
            p = (100 * torch.nn.functional.normalize(oc.encode_image(x), dim=-1)
                 @ torch.nn.functional.normalize(mine(prompts), dim=-1).T).softmax(-1)[0]
            print(f"[quiltnet/{act}] repo ADI-TCGA example tile zero-shot (true class adipose): top = {cls[p.argmax()]}"
                  f" p={p.max().item():.3f}")
            oc, mine = oc.to(dev), mine.to(dev)
            I = torch.cat([oc.encode_image(torch.stack([pre(x) for x in imgs[i:i + 64]]).to(dev)).cpu()
                           for i in range(0, len(imgs), 64)])
            T = torch.cat([mine(list(tn.caption[i:i + 64])).cpu() for i in range(0, len(tn), 64)])
            print(f"[quiltnet/{act}] retrieval (both towers {act}): {_retrieval(I, T)}")
        del oc, mine
        torch.cuda.empty_cache()


if __name__ == "__main__":
    if sys.argv[1:2] == ["verify"]:
        verify(sys.argv[2:] or ("keep", "quiltnet"))
    else:
        extract(sys.argv[1:])