File size: 4,114 Bytes
c2ede1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """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] # left padding -> last position is the last real token
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]) # length-sorted batches, restored below
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()
|