| """Emergent text alignment on ORION. The text head was trained ONLY on PathGen H&E-caption pairs (12_), the IF / |
| marker encoders ONLY on ORION H&E pairs (09_, --if_enc convnext_tiny --mk_enc mlp). Text and IF/markers never |
| co-occur in training; any text<->IF / text<->markers agreement is emergent through the shared UNI2 H&E space. |
| |
| (A) concept probes: prompts.json (PathGen-style sentences) scored against every ORION test tile embedding |
| (he / if / markers); ground truth = marker-positive fraction (or nuclei count). Metrics: Spearman rho and |
| AUROC (top quartile vs rest), on tissue tiles (>=20 nuclei; all tiles for 'cellularity'). |
| Baselines: zero-shot PLIP and ImageBind (their own image/text towers, H&E only, no training). |
| (B) retrieval with templated ORION descriptions (07_) as text queries through the PathGen-trained head |
| (domain-shifted: template style != PathGen captions), tie-aware R@1 on 100-tile galleries. |
| Outputs results/text_emergent_probes.csv, results/text_emergent_retrieval.csv |
| """ |
| import glob |
| import json |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn.functional as F |
| from scipy.stats import spearmanr |
| from sklearn.metrics import roc_auc_score |
|
|
| from common import CACHE, ROOT, RUNS, load_all_tiles, metrics |
|
|
| TXT = ["plip", "biomedclip", "imagebind"] |
| OUT = os.path.join(ROOT, "results") |
| prompts = json.load(open(os.path.join(ROOT, "prompts.json"))) |
| df = load_all_tiles() |
| stats = pd.read_csv(os.path.join(OUT, "tile_stats.csv")).set_index("tile_id") |
| df["n_nuclei"] = stats.loc[df.tile_id, "n_nuclei"].values |
|
|
|
|
| def load_head(enc, din, dout): |
| h = torch.nn.Sequential(torch.nn.Linear(din, 512), torch.nn.GELU(), torch.nn.Dropout(0.1), torch.nn.Linear(512, 512), |
| torch.nn.GELU(), torch.nn.LayerNorm(512), torch.nn.Linear(512, dout)) |
| h.load_state_dict(torch.load(os.path.join(RUNS, "text_pathgen", f"uni2_{enc}", "head.pt"), map_location="cpu")) |
| return h.eval() |
|
|
|
|
| def probe(scores, tiles, key, label_col): |
| sub = tiles if key == "cellularity" else tiles[df.loc[tiles, "n_nuclei"].values >= 20] |
| pos = np.isin(tiles, sub) |
| s = scores[pos] |
| y = df.loc[sub, label_col].values.astype(float) |
| if label_col != "n_nuclei": |
| y = y / np.maximum(df.loc[sub, "n_nuclei"].values, 1) |
| lab = y >= np.quantile(y, 0.75) |
| return {"spearman": float(spearmanr(s, y).correlation), |
| "auroc": float(roc_auc_score(lab, s)) if 0 < lab.sum() < len(lab) else np.nan, "n": int(len(sub))} |
|
|
|
|
| probe_rows, ret_rows = [], [] |
| for f in sorted(glob.glob(os.path.join(RUNS, "bench", "*_uni2_convnext_tiny_mlp_*", "test_embeddings.pt"))): |
| run = os.path.basename(os.path.dirname(f)) |
| split, fold = run.split("_fold")[0], int(run.split("_fold")[1].split("_")[0]) |
| E = torch.load(f, weights_only=False) |
| tiles, Z = np.asarray(E["tile_id"]), E["z"] |
| for enc in TXT: |
| prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{enc}.npy"))[:, 0].astype(np.float32)) |
| orion_txt = torch.from_numpy(np.load(os.path.join(CACHE, f"text_{enc}.npy"))[tiles, 0].astype(np.float32)) |
| head = load_head(enc, prm.shape[1], Z["he"].shape[1]) |
| with torch.no_grad(): |
| zp = F.normalize(head(prm), dim=-1) |
| zt = F.normalize(head(orion_txt), dim=-1) |
| for i, (key, pr) in enumerate(prompts.items()): |
| for m in ["he", "if", "markers"]: |
| sc = (Z[m] @ zp[i]).numpy() |
| probe_rows.append({"split": split, "fold": fold, "text_enc": enc, "concept": key, "modality": m, |
| "kind": "bound (PathGen text head)", **probe(sc, tiles, key, pr["label"])}) |
| for m in ["he", "if", "markers"]: |
| for q, g, name in [(zt, Z[m], f"text->{m}"), (Z[m], zt, f"{m}->text")]: |
| ret_rows.append({"split": split, "fold": fold, "text_enc": enc, "pair": name, **metrics(q, g)}) |
| |
| for fm in ["plip", "imagebind"]: |
| hz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"he_{fm}.npy"))[tiles, 0].astype(np.float32)), dim=-1) |
| pz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{fm}.npy"))[:, 0].astype(np.float32)), dim=-1) |
| for i, (key, pr) in enumerate(prompts.items()): |
| probe_rows.append({"split": split, "fold": fold, "text_enc": fm, "concept": key, "modality": "he", |
| "kind": "zero-shot (own image tower)", **probe((hz @ pz[i]).numpy(), tiles, key, pr["label"])}) |
|
|
| P = pd.DataFrame(probe_rows); R = pd.DataFrame(ret_rows) |
| P.to_csv(os.path.join(OUT, "text_emergent_probes.csv"), index=False) |
| R.to_csv(os.path.join(OUT, "text_emergent_retrieval.csv"), index=False) |
| pd.set_option("display.width", 250); pd.set_option("display.max_rows", 300) |
| print(P.groupby(["kind", "text_enc", "modality", "split"])[["auroc", "spearman"]].mean().round(3).unstack("split").to_string()) |
| print(P.groupby(["kind", "text_enc", "modality", "concept"])["auroc"].mean().unstack("concept").round(2).to_string()) |
| print(R.groupby(["text_enc", "pair", "split"])["R@1_g100"].mean().unstack("split").round(3).to_string()) |
|
|
