| """Frozen embeddings of every text set with additional (bio)medical / general text-embedding models. |
| |
| python 45_text_feats.py <name> [<name> ...] |
| |
| Text sets (same row order as the existing BiomedCLIP files): |
| alltxt cache/alltxt_meta.csv caption (TCGA real PathGen texts, 6480, junk rows kept, filtered later) |
| orion_gentext cache/orion_gentext_meta.csv caption (GPT synthetic captions of the ORION tiles, test only) |
| tcga_notext cache/tcga_notext_meta.csv caption (synthetic captions of TCGA images without real text, test only) |
| dlpfc_gentext cache/dlpfc_gentext_meta.csv caption (external DLPFC synthetic captions, test only) |
| Models (pooling / normalisation as on each model card; documents embedded without a query instruction): |
| biomedclip BiomedCLIP text tower (biomedclip_text.py, 256 tokens) - only for the sets that lack it |
| medcpt ncbi/MedCPT-Article-Encoder, [CLS], 512 tokens (PubMed search-log contrastive) |
| biolord FremyCompany/BioLORD-2023, mean pooling, 512 (clinical sentences / concepts, UMLS-based) |
| pubmedbert NeuML/pubmedbert-base-embeddings, mean pooling, 512 (PubMed title-abstract) |
| medembed abhinand/MedEmbed-large-v0.1 (bge-large medical fine-tune), [CLS], 512 |
| bge BAAI/bge-large-en-v1.5, [CLS], 512 (general-domain reference) |
| qwen3-0.6b / qwen3-4b Qwen/Qwen3-Embedding-*, last-token pooling, left padding, 1024 tokens |
| Output cache/<set>_text_<name>.npy (n, 1, D) float16 (raw, not normalised). |
| """ |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from transformers import AutoModel, AutoTokenizer |
|
|
| from common import CACHE |
|
|
| 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"} |
| HF = {"medcpt": ("ncbi/MedCPT-Article-Encoder", "cls", 512), "biolord": ("FremyCompany/BioLORD-2023", "mean", 512), |
| "pubmedbert": ("NeuML/pubmedbert-base-embeddings", "mean", 512), |
| "medembed": ("abhinand/MedEmbed-large-v0.1", "cls", 512), "bge": ("BAAI/bge-large-en-v1.5", "cls", 512), |
| "qwen3-0.6b": ("Qwen/Qwen3-Embedding-0.6B", "last", 1024), "qwen3-4b": ("Qwen/Qwen3-Embedding-4B", "last", 1024)} |
|
|
|
|
| def encoder(name): |
| if name == "biomedclip": |
| from biomedclip_text import BiomedCLIPText |
| m = BiomedCLIPText().to(dev).eval() |
| return m, 64 |
| repo, pool, L = HF[name] |
| tk = AutoTokenizer.from_pretrained(repo, padding_side="left" if pool == "last" else "right") |
| m = AutoModel.from_pretrained(repo, torch_dtype=torch.bfloat16 if pool == "last" else torch.float32).to(dev).eval() |
|
|
| def f(b): |
| x = tk(b, padding=True, truncation=True, max_length=L, return_tensors="pt").to(dev) |
| h = m(**x).last_hidden_state |
| if pool == "cls": |
| return h[:, 0] |
| if pool == "mean": |
| a = x["attention_mask"][..., None].to(h.dtype) |
| return (h * a).sum(1) / a.sum(1) |
| return h[:, -1] |
| return f, (16 if name == "qwen3-4b" else 64) |
|
|
|
|
| for name in sys.argv[1:]: |
| 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 = pd.read_csv(os.path.join(CACHE, meta)).caption.fillna("").astype(str).tolist() |
| order = np.argsort([len(t) for t in txt]) |
| E = None |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): |
| 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() |
|
|
