| """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) |
|
|
| 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): |
| import ftfy |
| 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")) |
| 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) |
| 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 |
|
|
| @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 |
|
|
|
|
| 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]) |
| 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() |
|
|
|
|
| |
| 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 |
| 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 |
| 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) |
| 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:]) |
|
|
