#!/usr/bin/env python3 """ Novel Embedding Probes v2 — GlycanBERT V5 ========================================== Uses the FULL pretraining data (254k WURCS) instead of benchmark subsets. All plots use Nature BGP color palette, 300 DPI, publication-ready. Probes: 1. Ambiguity (? marks) — 98k ambiguous vs 156k clean WURCS 2. Composition — monosaccharide fingerprint from [CLS] 3. KNN Purity (expanded) — domain, kingdom, link (N vs O), immunogenicity 4. Polymerization — chain length / branch depth regression 5. Size Prediction — small/med/large/xlarge from frozen [CLS] 6. N-vs-O Link (binary) — only N and O linkages embedded 7. MLM Zero-Shot (fixed) — random token replacement instead of [MASK] 8. Token Importance (fixed) — leave-one-out CLS shift analysis Usage: python novel_probes_v2.py --model v5 --probe all --max_samples 5000 """ import os, sys, json, argparse, csv import numpy as np from pathlib import Path from collections import Counter # ─── Paths ─────────────────────────────────────────────────────────────── PROJECT_ROOT = Path(__file__).resolve().parents[2] VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json' CHECKPOINTS = { 'v5': PROJECT_ROOT / 'checkpoints_v5_bpe_topo' / 'best_v5_bpe_topo_model.pt', 'v6': PROJECT_ROOT / 'bert_v5.1_contrastive' / 'checkpoints' / 'best_v51_contrastive_model.pt', } PRETRAIN_CSV = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'multimodal_index.csv' BENCH_DIR = PROJECT_ROOT / 'bench' / 'GlycanML' / 'data' # ─── Nature BGP Color Palette ────────────────────────────────────────── # From: https://www.nature.com/documents/natrev-artworkguide.pdf NATURE_COLORS = { 'blue': '#0072B2', 'orange': '#E69F00', 'green': '#009E73', 'red': '#D55E00', 'purple': '#CC79A7', 'cyan': '#56B4E9', 'yellow': '#F0E442', 'black': '#000000', 'grey': '#999999', } # Categorical palettes PALETTE_2 = ['#0072B2', '#D55E00'] PALETTE_3 = ['#0072B2', '#E69F00', '#009E73'] PALETTE_4 = ['#0072B2', '#E69F00', '#009E73', '#D55E00'] PALETTE_5 = ['#0072B2', '#E69F00', '#009E73', '#D55E00', '#CC79A7'] PALETTE_8 = ['#0072B2', '#E69F00', '#009E73', '#D55E00', '#CC79A7', '#56B4E9', '#F0E442', '#999999'] PALETTE_11 = PALETTE_8 + ['#000000', '#882255', '#44AA99'] def get_palette(n): if n <= 2: return PALETTE_2[:n] if n <= 3: return PALETTE_3[:n] if n <= 4: return PALETTE_4[:n] if n <= 5: return PALETTE_5[:n] if n <= 8: return PALETTE_8[:n] return (PALETTE_11 * ((n // 11) + 1))[:n] # ─── Plot setup ───────────────────────────────────────────────────────── def setup_nature_style(): import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt plt.rcParams.update({ 'font.family': 'sans-serif', 'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans'], 'font.size': 10, 'axes.titlesize': 12, 'axes.labelsize': 11, 'xtick.labelsize': 9, 'ytick.labelsize': 9, 'legend.fontsize': 9, 'figure.dpi': 300, 'savefig.dpi': 300, 'savefig.bbox': 'tight', 'axes.linewidth': 0.8, 'axes.spines.top': False, 'axes.spines.right': False, }) return plt # ─── Model loading ────────────────────────────────────────────────────── # Matches the working pattern from embed_benchmark_tasks.py and extract_embeddings.py sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(PROJECT_ROOT / 'bert_training_v4')) from model.multimodal_glycan_bert_v3 import MultimodalGlycanBERT, MultimodalGlycanBERTConfig from downstream_tasks.utils.tokenizer import WURCSTokenizer def load_model(ckpt_path, device='cuda'): """Load MultimodalGlycanBERT from checkpoint (matches embed_benchmark_tasks.py).""" import torch print(f"Loading model from {ckpt_path}...") ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False) if 'model_state_dict' in ckpt: state_dict = ckpt['model_state_dict'] else: state_dict = ckpt # Strip projection head keys (V6 only) backbone_sd = {k: v for k, v in state_dict.items() if not k.startswith('proj_head.')} n_stripped = len(state_dict) - len(backbone_sd) if n_stripped > 0: print(f" Stripped {n_stripped} projection head keys") # Infer vocab size from state dict vocab_size = backbone_sd['seq_embeddings.token_embeddings.weight'].shape[0] # Check for MS embeddings ms_total_vocab = None if 'ms_embeddings.token_embeddings.weight' in backbone_sd: ms_total_vocab = backbone_sd['ms_embeddings.token_embeddings.weight'].shape[0] config_kwargs = dict( seq_vocab_size=vocab_size, seq_hidden_size=768, seq_num_layers=12, seq_num_heads=12, seq_max_length=256, use_cnn_frontend=True, cnn_kernel_size=3, ) if ms_total_vocab is not None: config_kwargs['ms_vocab_size'] = ms_total_vocab - vocab_size config = MultimodalGlycanBERTConfig(**config_kwargs) model = MultimodalGlycanBERT(config) model.load_state_dict(backbone_sd, strict=False) model.to(device) model.eval() n_params = sum(p.numel() for p in model.parameters()) print(f" Model loaded: {n_params:,} params, vocab_size={vocab_size}") return model # ─── Data loading ─────────────────────────────────────────────────────── def load_pretrain_wurcs(tokenizer, max_n=None): """Load ALL WURCS from multimodal_index.csv + metadata.""" samples = [] with open(PRETRAIN_CSV) as f: reader = csv.DictReader(f) for row in reader: w = row['wurcs'] if not w.startswith('WURCS'): continue try: n_res = int(w.split('/')[1].split(',')[1]) if '/' in w else 0 except: n_res = 0 has_q = '?' in w q_count = w.count('?') samples.append({ 'wurcs': w, 'accession': row.get('accession', ''), 'n_residues': n_res, 'has_ambiguity': has_q, 'ambiguity_count': q_count, 'monosaccharide_names': row.get('monosaccharide_names', ''), }) if max_n and len(samples) >= max_n: break print(f" Loaded {len(samples)} WURCS from pretraining data") print(f" Ambiguous (has ?): {sum(1 for s in samples if s['has_ambiguity'])}") return samples def load_benchmark_glycans(tokenizer, csv_name, max_n=None): """Load glycans from a benchmark CSV.""" csv_path = BENCH_DIR / csv_name if not csv_path.exists(): print(f" WARNING: {csv_path} not found") return [] samples = [] with open(csv_path) as f: reader = csv.DictReader(f) for row in reader: w = row.get('wurcs', '') if not w.startswith('WURCS'): continue samples.append(row) if max_n and len(samples) >= max_n: break return samples # ─── Embedding ────────────────────────────────────────────────────────── # Matches extract_embeddings.py pattern: use model.seq_embeddings() with # branch_depths and linkage_types from the WURCSTokenizer. def batch_cls_embeddings(model, samples, device='cuda', batch_size=64, max_len=256): """Extract [CLS] embeddings for a list of samples. Uses WURCSTokenizer.tokenize() to get token_ids, branch_depths, and linkage_types, then runs model.seq_embeddings() — the working forward pass pattern from extract_embeddings.py. """ import torch import torch.nn.functional as F tokenizer = WURCSTokenizer(str(VOCAB_PATH)) if not samples: return np.zeros((0, 768)) all_embs = [] n_errors = 0 for i in range(0, len(samples), batch_size): batch = samples[i:i+batch_size] batch_embs = [] for s in batch: try: result = tokenizer.tokenize(s['wurcs'], max_length=max_len) token_ids = torch.tensor(result['token_ids'], dtype=torch.long) branch_depths = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long) linkage_types = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long) # Ensure same length min_l = min(len(token_ids), len(branch_depths), len(linkage_types)) token_ids = token_ids[:min_l] branch_depths = branch_depths[:min_l] linkage_types = linkage_types[:min_l] # Truncate / pad to max_len if min_l > max_len: token_ids = token_ids[:max_len] branch_depths = branch_depths[:max_len] linkage_types = linkage_types[:max_len] elif min_l < max_len: pad_len = max_len - min_l token_ids = F.pad(token_ids, (0, pad_len), value=0) branch_depths = F.pad(branch_depths, (0, pad_len), value=0) linkage_types = F.pad(linkage_types, (0, pad_len), value=0) # Forward through seq encoder token_ids = token_ids.unsqueeze(0).to(device) branch_depths = branch_depths.unsqueeze(0).to(device) linkage_types = linkage_types.unsqueeze(0).to(device) with torch.no_grad(): seq_out = model.seq_embeddings(token_ids, branch_depths=branch_depths, linkage_types=linkage_types) cls_emb = seq_out[0, 0, :].cpu().numpy() batch_embs.append(cls_emb) except Exception as e: n_errors += 1 if n_errors <= 5: import traceback as tb print(f" ERROR (sample {i}): {e}") tb.print_exc() batch_embs.append(np.zeros(768)) all_embs.extend(batch_embs) if (i // batch_size) % 20 == 0 and i > 0: print(f" Embedded {i}/{len(samples)} ({n_errors} errors)") if n_errors > 0: print(f" WARNING: {n_errors}/{len(samples)} tokenization errors") print(f" Embedded {len(all_embs)} total samples", flush=True) return np.array(all_embs) if all_embs else np.zeros((0, 768)) # ═══════════════════════════════════════════════════════════════════════ # PROBE 1: Ambiguity (? marks) — FULL DATA # ═══════════════════════════════════════════════════════════════════════ def save_publication_plots(X, labels, label_name, out_dir, title_prefix=""): """Generate PCA publication-quality plot (UMAP disabled to save memory).""" import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from sklearn.decomposition import PCA out_dir = Path(out_dir) safe_name = re.sub(r'[^a-zA-Z0-9_-]', '_', label_name.lower().replace(' ', '_')) # Subsample for plotting max_plot = min(len(X), 10000) if len(X) > max_plot: idx = np.random.RandomState(42).choice(len(X), max_plot, replace=False) X_sub = X[idx] labels_sub = [labels[i] for i in idx] else: X_sub = X labels_sub = list(labels) unique_labels = sorted(set(labels_sub)) cmap = plt.cm.get_cmap('tab20', len(unique_labels)) color_map = {lbl: cmap(i) for i, lbl in enumerate(unique_labels)} fig, ax = plt.subplots(1, 1, figsize=(8, 6)) pca = PCA(n_components=2) X_pca = pca.fit_transform(X_sub) for lbl in unique_labels: mask = [l == lbl for l in labels_sub] ax.scatter(X_pca[np.array(mask), 0], X_pca[np.array(mask), 1], c=[color_map[lbl]], label=lbl, s=8, alpha=0.5) ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})') ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})') ax.set_title(f'{title_prefix} — PCA by {label_name}') ax.legend(fontsize=6, markerscale=2, loc='best') plt.tight_layout() plt.savefig(out_dir / f'{safe_name}_pca.png', dpi=150) plt.close(fig) del X_sub, X_pca, fig print(f" Saved: {safe_name}_pca.png") def probe_ambiguity(model, tokenizer, device, output_dir, max_samples=10000, **kwargs): print("\n" + "="*60) print("PROBE 1: Ambiguity Analysis (? marks in WURCS)") print("="*60) samples = load_pretrain_wurcs(tokenizer, max_n=max_samples) ambig = [s for s in samples if s['has_ambiguity']] clean = [s for s in samples if not s['has_ambiguity']] print(f" Ambiguous: {len(ambig)}, Clean: {len(clean)}") # Bin ambiguity into levels for s in samples: qc = s['ambiguity_count'] if qc == 0: s['amb_level'] = 'none' elif qc <= 2: s['amb_level'] = 'low (1-2)' elif qc <= 5: s['amb_level'] = 'medium (3-5)' else: s['amb_level'] = 'high (6+)' level_counts = Counter(s['amb_level'] for s in samples) print(f" Levels: {dict(level_counts)}") # Subsample for balance if needed min_group = min(len(ambig), len(clean), 2000) np.random.seed(42) if len(ambig) > min_group: ambig = [ambig[i] for i in np.random.choice(len(ambig), min_group, replace=False)] if len(clean) > min_group: clean = [clean[i] for i in np.random.choice(len(clean), min_group, replace=False)] all_samp = ambig + clean print(f" Embedding {len(all_samp)} samples (balanced)...") embeddings = batch_cls_embeddings(model, all_samp, device=device) if embeddings.shape[0] == 0: print(" SKIPPING probe_ambiguity — no valid embeddings") return {'error': 'no_embeddings', 'n_ambig': len(ambig), 'n_clean': len(clean)} labels = ['ambiguous']*len(ambig) + ['clean']*len(clean) # Metrics from sklearn.metrics import silhouette_score from sklearn.neighbors import KNeighborsClassifier from sklearn.model_selection import cross_val_score from sklearn.preprocessing import StandardScaler int_labels = np.array([0 if l == 'ambiguous' else 1 for l in labels]) sil = float(silhouette_score(embeddings, int_labels)) X = StandardScaler().fit_transform(embeddings) knn_acc = float(cross_val_score( KNeighborsClassifier(n_neighbors=10), X, int_labels, cv=5, scoring='accuracy' ).mean()) print(f" Silhouette (ambig vs clean): {sil:.4f}") print(f" KNN classification accuracy: {knn_acc:.4f}") # Cosine similarity analysis from sklearn.metrics.pairwise import cosine_similarity emb_ambig = embeddings[:len(ambig)] emb_clean = embeddings[len(ambig):] within_ambig = float(np.mean(cosine_similarity(emb_ambig))) within_clean = float(np.mean(cosine_similarity(emb_clean))) between = float(np.mean(cosine_similarity(emb_ambig, emb_clean))) print(f" Within-ambig sim: {within_ambig:.4f}") print(f" Within-clean sim: {within_clean:.4f}") print(f" Between sim: {between:.4f}") # t-SNE plot plt = setup_nature_style() from sklearn.manifold import TSNE perp = min(30, len(embeddings) - 1) coords = TSNE(n_components=2, perplexity=perp, max_iter=1000, init='pca', random_state=42, learning_rate='auto').fit_transform(embeddings) fig, ax = plt.subplots(figsize=(8, 6)) colors = {'clean': NATURE_COLORS['blue'], 'ambiguous': NATURE_COLORS['orange']} for label in ['clean', 'ambiguous']: mask = np.array(labels) == label ax.scatter(coords[mask, 0], coords[mask, 1], c=colors[label], label=f'{label} (n={mask.sum()})', s=8, alpha=0.5, edgecolors='none') ax.set_title(f'Ambiguity Probe: WURCS with ? marks vs Clean\n' f'Silhouette={sil:.4f} | KNN Acc={knn_acc:.4f}') ax.set_xlabel('t-SNE 1') ax.set_ylabel('t-SNE 2') ax.legend(loc='best', framealpha=0.8) plt.tight_layout() plt.savefig(os.path.join(output_dir, 'ambiguity_probe.png'), dpi=300, bbox_inches='tight') plt.close() results = { 'silhouette': sil, 'knn_accuracy': knn_acc, 'within_ambig_sim': within_ambig, 'within_clean_sim': within_clean, 'between_sim': between, 'n_ambiguous': len(ambig), 'n_clean': len(clean), 'level_counts': dict(level_counts), } with open(os.path.join(output_dir, 'ambiguity_probe.json'), 'w') as f: json.dump(results, f, indent=2, default=str) return results # ═══════════════════════════════════════════════════════════════════════ # PROBE 2: Monosaccharide Composition — FULL DATA # ═══════════════════════════════════════════════════════════════════════ def probe_composition(model, tokenizer, device, output_dir, max_samples=5000, **kwargs): _cached_embs = kwargs.get("_cached_embs") _cached_samples = kwargs.get("_cached_samples") print("\n" + "="*60) print("PROBE 2: Monosaccharide Composition") print("="*60) if _cached_samples is not None: samples = _cached_samples[:max_samples] else: samples = load_pretrain_wurcs(tokenizer, max_n=max_samples) # Parse monosaccharide names for s in samples: names = s.get('monosaccharide_names', '') s['monos'] = [m.strip() for m in names.split(',') if m.strip()] if names else [] # Find top-20 most common monosaccharides all_monos = [] for s in samples: all_monos.extend(s['monos']) mono_counts = Counter(all_monos) top_k = [m for m, _ in mono_counts.most_common(20)] print(f" Top-20 monos: {top_k[:5]}...") embeddings = batch_cls_embeddings(model, samples, device=device) from sklearn.linear_model import LogisticRegression from sklearn.model_selection import cross_val_score from sklearn.preprocessing import StandardScaler X = StandardScaler().fit_transform(embeddings) per_mono_results = {} for mono in top_k: y = np.array([1 if mono in s['monos'] else 0 for s in samples]) n_pos = int(y.sum()) if n_pos < 20 or n_pos > len(y) - 20: continue scores = cross_val_score( LogisticRegression(max_iter=500, class_weight='balanced'), X, y, cv=5, scoring='roc_auc' ) per_mono_results[mono] = {'auc': float(scores.mean()), 'std': float(scores.std()), 'n_pos': n_pos} print(f" {mono:35s}: AUC={scores.mean():.4f} ± {scores.std():.4f} (n+={n_pos})") # Bar chart plt = setup_nature_style() monos_sorted = sorted(per_mono_results.keys(), key=lambda m: per_mono_results[m]['auc'], reverse=True) fig, ax = plt.subplots(figsize=(12, 6)) x = range(len(monos_sorted)) aucs = [per_mono_results[m]['auc'] for m in monos_sorted] stds = [per_mono_results[m]['std'] for m in monos_sorted] bars = ax.bar(x, aucs, yerr=stds, color=NATURE_COLORS['blue'], alpha=0.8, edgecolor='white', linewidth=0.5, capsize=3) ax.axhline(0.5, color=NATURE_COLORS['grey'], linestyle='--', linewidth=0.8, label='Random baseline') ax.set_xticks(x) ax.set_xticklabels(monos_sorted, rotation=45, ha='right', fontsize=7) ax.set_ylabel('ROC AUC') ax.set_title(f'Monosaccharide Detection from Frozen [CLS] Embedding\n(n={len(samples)}, {len(monos_sorted)} monosaccharides)') ax.set_ylim(0.4, 1.05) ax.legend(loc='lower right') plt.tight_layout() plt.savefig(os.path.join(output_dir, 'composition_probe.png'), dpi=300, bbox_inches='tight') plt.close() results = {'per_mono_auc': per_mono_results, 'n_samples': len(samples), 'top_k': top_k} with open(os.path.join(output_dir, 'composition_probe.json'), 'w') as f: json.dump(results, f, indent=2, default=str) # Publication UMAP — color by top-3 monosaccharides try: from pathlib import Path top3 = top_k[:3] labels_mono = [] for s in samples[:len(embs)]: monos = set(m.strip() for m in s.get('monosaccharide_names', '').split(',')) found = [m for m in top3 if m in monos] labels_mono.append(found[0] if len(found) == 1 else ('Multi' if len(found) > 1 else 'None')) save_publication_plots(np.array(embs), labels_mono, 'Top-3 Monosaccharides', Path(args.output_dir), title_prefix='Probe 2') except Exception as e: print(f" Pub plot error: {e}") return results def probe_knn_purity(model, tokenizer, device, output_dir, max_samples=15000, **kwargs): print("\n" + "="*60) print("PROBE 3: KNN Purity (Expanded)") print("="*60) # Load classification data (domain + kingdom) cls_samples = load_benchmark_glycans(tokenizer, 'glycan_classification_wurcs_subset.csv', max_n=max_samples) # Load link data link_samples = load_benchmark_glycans(tokenizer, 'glycan_link_wurcs_subset.csv', max_n=max_samples) # Load immunogenicity immuno_samples = load_benchmark_glycans(tokenizer, 'glycan_immunogenicity_wurcs_subset.csv', max_n=max_samples) results = {} # Domain + Kingdom KNN if cls_samples: print(f" Classification samples: {len(cls_samples)}") cls_embs = batch_cls_embeddings(model, cls_samples, device=device) for task_col in ['domain', 'kingdom']: labels = [] valid_mask = [] for i, s in enumerate(cls_samples): label = s.get(task_col, '') if label: labels.append(label) valid_mask.append(i) if not labels: continue embs = cls_embs[valid_mask] label_arr = np.array(labels) n_classes = len(set(labels)) class_counts = Counter(labels) print(f" {task_col}: {len(labels)} samples, {n_classes} classes") print(f" Distribution: {dict(class_counts)}") for k in [5, 10, 20, 50]: from sklearn.neighbors import NearestNeighbors nn = NearestNeighbors(n_neighbors=k+1, metric='cosine') nn.fit(embs) _, indices = nn.kneighbors(embs) purities = [] for i in range(len(embs)): neighbors = indices[i, 1:] # exclude self same_class = np.sum(label_arr[neighbors] == label_arr[i]) purities.append(same_class / k) purity = float(np.mean(purities)) results[f'{task_col}_k{k}'] = purity print(f" KNN Purity (k={k:2d}): {purity:.4f}") # Per-class purity at k=10 for cls_name in sorted(set(labels)): cls_mask = label_arr == cls_name cls_purity = float(np.mean([purities[i] for i in range(len(purities)) if cls_mask[i]])) results[f'{task_col}_{cls_name}_k10'] = cls_purity # Link KNN (N vs O only — binary) if link_samples: no_samples = [s for s in link_samples if s.get('target', '') in ('N', 'O')] print(f" Link N-vs-O samples: {len(no_samples)}") if len(no_samples) > 50: link_embs = batch_cls_embeddings(model, no_samples, device=device) link_labels = np.array([s['target'] for s in no_samples]) for k in [5, 10, 20]: from sklearn.neighbors import NearestNeighbors nn = NearestNeighbors(n_neighbors=k+1, metric='cosine') nn.fit(link_embs) _, indices = nn.kneighbors(link_embs) purities = [] for i in range(len(link_embs)): neighbors = indices[i, 1:] same_class = np.sum(link_labels[neighbors] == link_labels[i]) purities.append(same_class / k) purity = float(np.mean(purities)) results[f'link_NO_k{k}'] = purity print(f" Link N-vs-O KNN (k={k:2d}): {purity:.4f}") # Immunogenicity KNN if immuno_samples: print(f" Immunogenicity samples: {len(immuno_samples)}") i_embs = batch_cls_embeddings(model, immuno_samples, device=device) i_labels = np.array([s.get('target', s.get('immunogenicity', '')) for s in immuno_samples]) for k in [5, 10, 20]: from sklearn.neighbors import NearestNeighbors nn = NearestNeighbors(n_neighbors=k+1, metric='cosine') nn.fit(i_embs) _, indices = nn.kneighbors(i_embs) purities = [] for i in range(len(i_embs)): neighbors = indices[i, 1:] same = np.sum(i_labels[neighbors] == i_labels[i]) purities.append(same / k) purity = float(np.mean(purities)) results[f'immunogenicity_k{k}'] = purity print(f" Immunogenicity KNN (k={k:2d}): {purity:.4f}") with open(os.path.join(output_dir, 'knn_purity.json'), 'w') as f: json.dump(results, f, indent=2, default=str) # Publication UMAP — color by domain try: from pathlib import Path domain_labels_all = [s.get('domain','') for s in samples[:len(embs)]] if len(set(l for l in domain_labels_all if l)) >= 2: save_publication_plots(np.array(embs), domain_labels_all, 'Taxonomy Domain', Path(args.output_dir), title_prefix='Probe 3a') except Exception as e: print(f" Pub plot error: {e}") return results def probe_polymerization(model, tokenizer, device, output_dir, max_samples=5000, **kwargs): _cached_embs = kwargs.get("_cached_embs") _cached_samples = kwargs.get("_cached_samples") print("\n" + "="*60) print("PROBE 4: Polymerization / Complexity Probe") print("="*60) if _cached_samples is not None: samples = _cached_samples[:max_samples] else: samples = load_pretrain_wurcs(tokenizer, max_n=max_samples) # Parse complexity features from WURCS for s in samples: w = s['wurcs'] try: parts = w.split('/') counts = parts[1].split(',') s['n_unique_res'] = int(counts[0]) s['n_total_res'] = int(counts[1]) s['n_linkages'] = int(counts[2]) if len(counts) > 2 else s['n_total_res'] - 1 except: s['n_unique_res'] = s['n_residues'] s['n_total_res'] = s['n_residues'] s['n_linkages'] = max(0, s['n_residues'] - 1) # Branch depth try: link_str = w.split('/')[-1] if '/' in w else '' depth = link_str.count('-') - (s['n_total_res'] - 1) if link_str else 0 s['branch_depth'] = max(0, depth) except: s['branch_depth'] = 0 print(f" Samples: {len(samples)}, Residues: {min(s['n_total_res'] for s in samples)}-{max(s['n_total_res'] for s in samples)}") embeddings = batch_cls_embeddings(model, samples, device=device) from sklearn.linear_model import Ridge from sklearn.model_selection import cross_val_score from sklearn.preprocessing import StandardScaler from scipy.stats import spearmanr X = StandardScaler().fit_transform(embeddings) features = { 'n_total_residues': [s['n_total_res'] for s in samples], 'n_unique_residues': [s['n_unique_res'] for s in samples], 'n_linkages': [s['n_linkages'] for s in samples], 'branch_depth': [s['branch_depth'] for s in samples], } results = {'linear_probe_r2': {}, 'spearman_correlations': {}, 'n_samples': len(samples)} for fname, values in features.items(): y = np.array(values, dtype=float) if np.std(y) < 1e-6: continue r2 = cross_val_score(Ridge(alpha=1.0), X, y, cv=5, scoring='r2') results['linear_probe_r2'][fname] = {'r2_mean': float(r2.mean()), 'r2_std': float(r2.std())} print(f" {fname:25s}: R²={r2.mean():.4f} ± {r2.std():.4f}") # Pairwise Spearman from sklearn.metrics.pairwise import euclidean_distances dists = euclidean_distances(embeddings) upper_idx = np.triu_indices(len(embeddings), k=1) emb_dists = dists[upper_idx] for fname, values in features.items(): y = np.array(values, dtype=float) feat_diffs = np.abs(y[upper_idx[0]] - y[upper_idx[1]]) # Subsample for speed if len(emb_dists) > 500000: idx = np.random.choice(len(emb_dists), 500000, replace=False) rho, p = spearmanr(emb_dists[idx], feat_diffs[idx]) else: rho, p = spearmanr(emb_dists, feat_diffs) results['spearman_correlations'][fname] = {'rho': float(rho), 'p': float(p)} print(f" Spearman ρ ({fname}): {rho:.4f} (p={p:.2e})") # Scatter plot: n_residues vs CLS PCA1 plt = setup_nature_style() from sklearn.decomposition import PCA pca = PCA(n_components=2) pca_coords = pca.fit_transform(embeddings) fig, axes = plt.subplots(1, 2, figsize=(14, 5)) for ax, (fname, values) in zip(axes, [('n_total_residues', features['n_total_residues']), ('n_unique_residues', features['n_unique_residues'])]): sc = ax.scatter(pca_coords[:, 0], pca_coords[:, 1], c=values, cmap='viridis', s=5, alpha=0.5, edgecolors='none') plt.colorbar(sc, ax=ax, label=fname) ax.set_xlabel('PCA 1') ax.set_ylabel('PCA 2') ax.set_title(f'{fname}\nR²={results["linear_probe_r2"].get(fname, {}).get("r2_mean", 0):.4f}') plt.suptitle(f'Polymerization Probe (n={len(samples)})', fontsize=13) plt.tight_layout() plt.savefig(os.path.join(output_dir, 'polymerization_probe.png'), dpi=300, bbox_inches='tight') plt.close() with open(os.path.join(output_dir, 'polymerization_probe.json'), 'w') as f: json.dump(results, f, indent=2, default=str) return results # ═══════════════════════════════════════════════════════════════════════ # PROBE 5: Size Category Prediction — FULL DATA # ═══════════════════════════════════════════════════════════════════════ def probe_size(model, tokenizer, device, output_dir, max_samples=5000, **kwargs): _cached_embs = kwargs.get("_cached_embs") _cached_samples = kwargs.get("_cached_samples") print("\n" + "="*60) print("PROBE 5: Size Category Prediction") print("="*60) if _cached_samples is not None: samples = _cached_samples[:max_samples] else: samples = load_pretrain_wurcs(tokenizer, max_n=max_samples) for s in samples: n = s['n_residues'] s['size'] = 'small' if n <= 3 else ('medium' if n <= 6 else ('large' if n <= 10 else 'very_large')) size_dist = Counter(s['size'] for s in samples) print(f" Sizes: {dict(size_dist)}") embeddings = batch_cls_embeddings(model, samples, device=device) labels = [s['size'] for s in samples] from sklearn.metrics import silhouette_score from sklearn.neighbors import KNeighborsClassifier from sklearn.model_selection import cross_val_score from sklearn.preprocessing import StandardScaler unique = sorted(set(labels)) l2i = {l: i for i, l in enumerate(unique)} int_labels = np.array([l2i[l] for l in labels]) sil = float(silhouette_score(embeddings, int_labels)) X = StandardScaler().fit_transform(embeddings) knn_acc = float(cross_val_score( KNeighborsClassifier(n_neighbors=10), X, int_labels, cv=5, scoring='accuracy' ).mean()) print(f" Silhouette: {sil:.4f}, KNN Acc: {knn_acc:.4f}") # t-SNE plot plt = setup_nature_style() from sklearn.manifold import TSNE perp = min(30, len(embeddings) - 1) coords = TSNE(n_components=2, perplexity=perp, max_iter=1000, init='pca', random_state=42, learning_rate='auto').fit_transform(embeddings) fig, ax = plt.subplots(figsize=(8, 6)) colors = {'small': NATURE_COLORS['green'], 'medium': NATURE_COLORS['blue'], 'large': NATURE_COLORS['orange'], 'very_large': NATURE_COLORS['red']} for cat in ['small', 'medium', 'large', 'very_large']: mask = np.array(labels) == cat if mask.any(): ax.scatter(coords[mask, 0], coords[mask, 1], c=colors[cat], label=f'{cat} (n={mask.sum()})', s=8, alpha=0.5, edgecolors='none') ax.set_title(f'Size Category Prediction\nSilhouette={sil:.4f} | KNN Acc={knn_acc:.4f}') ax.set_xlabel('t-SNE 1') ax.set_ylabel('t-SNE 2') ax.legend(loc='best', framealpha=0.8) plt.tight_layout() plt.savefig(os.path.join(output_dir, 'size_probe.png'), dpi=300, bbox_inches='tight') plt.close() results = {'silhouette': sil, 'knn_accuracy': knn_acc, 'sizes': dict(size_dist)} with open(os.path.join(output_dir, 'size_probe.json'), 'w') as f: json.dump(results, f, indent=2, default=str) return results # ═══════════════════════════════════════════════════════════════════════ # PROBE 6: N-vs-O Link (Binary Embedding) # ═══════════════════════════════════════════════════════════════════════ def probe_link_binary(model, tokenizer, device, output_dir, max_samples=5000, **kwargs): print("\n" + "="*60) print("PROBE 6: N-linked vs O-linked (Binary)") print("="*60) link_samples = load_benchmark_glycans(tokenizer, 'glycan_link_wurcs_subset.csv', max_n=max_samples) # Filter to N and O only no_samples = [s for s in link_samples if s.get('target', '') in ('N', 'O')] print(f" N+O samples: {len(no_samples)}") label_dist = Counter(s['target'] for s in no_samples) print(f" Distribution: {dict(label_dist)}") if len(no_samples) < 50: print(" Too few samples, skipping") return {'error': 'Too few N/O samples'} embeddings = batch_cls_embeddings(model, no_samples, device=device) labels = [s['target'] for s in no_samples] from sklearn.metrics import silhouette_score from sklearn.neighbors import KNeighborsClassifier from sklearn.model_selection import cross_val_score from sklearn.preprocessing import StandardScaler from sklearn.linear_model import LogisticRegression int_labels = np.array([0 if l == 'N' else 1 for l in labels]) sil = float(silhouette_score(embeddings, int_labels)) X = StandardScaler().fit_transform(embeddings) knn_acc = float(cross_val_score( KNeighborsClassifier(n_neighbors=10), X, int_labels, cv=5, scoring='accuracy' ).mean()) lr_auc = float(cross_val_score( LogisticRegression(max_iter=500), X, int_labels, cv=5, scoring='roc_auc' ).mean()) print(f" Silhouette: {sil:.4f}, KNN Acc: {knn_acc:.4f}, LR AUC: {lr_auc:.4f}") # t-SNE plt = setup_nature_style() from sklearn.manifold import TSNE perp = min(30, len(embeddings) - 1) coords = TSNE(n_components=2, perplexity=perp, max_iter=1000, init='pca', random_state=42, learning_rate='auto').fit_transform(embeddings) fig, ax = plt.subplots(figsize=(8, 6)) for label, color, name in [('N', NATURE_COLORS['blue'], 'N-linked'), ('O', NATURE_COLORS['orange'], 'O-linked')]: mask = np.array(labels) == label ax.scatter(coords[mask, 0], coords[mask, 1], c=color, label=f'{name} (n={mask.sum()})', s=15, alpha=0.6, edgecolors='none') ax.set_title(f'N-linked vs O-linked Glycans\nSilhouette={sil:.4f} | KNN={knn_acc:.4f} | AUC={lr_auc:.4f}') ax.set_xlabel('t-SNE 1') ax.set_ylabel('t-SNE 2') ax.legend(loc='best', framealpha=0.8) plt.tight_layout() plt.savefig(os.path.join(output_dir, 'link_binary_probe.png'), dpi=300, bbox_inches='tight') plt.close() results = {'silhouette': sil, 'knn_accuracy': knn_acc, 'lr_auc': lr_auc, 'distribution': dict(label_dist)} with open(os.path.join(output_dir, 'link_binary_probe.json'), 'w') as f: json.dump(results, f, indent=2, default=str) return results # ═══════════════════════════════════════════════════════════════════════ # PROBE 7: MLM Zero-Shot (Fixed — random token replacement) # ═══════════════════════════════════════════════════════════════════════ def probe_mlm_zeroshot(model, tokenizer, device, output_dir, max_samples=500, **kwargs): _cached_embs = kwargs.get("_cached_embs") _cached_samples = kwargs.get("_cached_samples") print("\n" + "="*60) print("PROBE 7: MLM Zero-Shot (Token Replacement)") print("="*60) import torch import torch.nn.functional as F if _cached_samples is not None: samples = _cached_samples[:max_samples] else: samples = load_pretrain_wurcs(tokenizer, max_n=max_samples) tok = WURCSTokenizer(str(VOCAB_PATH)) MAX_LEN = 256 def _get_cls(token_ids, branch_depths, linkage_types): """Helper: pad/truncate and run model.seq_embeddings(), return CLS numpy.""" min_l = min(len(token_ids), len(branch_depths), len(linkage_types)) token_ids = token_ids[:min_l] branch_depths = branch_depths[:min_l] linkage_types = linkage_types[:min_l] if min_l > MAX_LEN: token_ids = token_ids[:MAX_LEN] branch_depths = branch_depths[:MAX_LEN] linkage_types = linkage_types[:MAX_LEN] elif min_l < MAX_LEN: p = MAX_LEN - min_l token_ids = F.pad(token_ids, (0, p), value=0) branch_depths = F.pad(branch_depths, (0, p), value=0) linkage_types = F.pad(linkage_types, (0, p), value=0) with torch.no_grad(): out = model.seq_embeddings( token_ids.unsqueeze(0).to(device), branch_depths=branch_depths.unsqueeze(0).to(device), linkage_types=linkage_types.unsqueeze(0).to(device), ) return out[0, 0, :].cpu().numpy() correct_predictions = 0 total_predictions = 0 per_position_shifts = [] for idx, s in enumerate(samples): if idx > 200: break # cap for speed try: result = tok.tokenize(s['wurcs'], max_length=MAX_LEN) ids = torch.tensor(result['token_ids'], dtype=torch.long) bd = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long) lt = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long) real_len = result.get('length', len(ids)) if real_len < 3: continue # Get original CLS cls_orig = _get_cls(ids.clone(), bd.clone(), lt.clone()) # For each non-special token position, replace with UNK for pos in range(1, min(real_len - 1, 20)): original_token = ids[pos].item() ids_modified = ids.clone() ids_modified[pos] = 1 # UNK token cls_modified = _get_cls(ids_modified, bd.clone(), lt.clone()) shift = float(np.linalg.norm(cls_orig - cls_modified)) per_position_shifts.append({ 'sample_idx': idx, 'position': pos, 'original_token': original_token, 'cls_shift': shift, }) total_predictions += 1 except Exception as e: continue if idx % 50 == 0: print(f" Processed {idx}/{min(len(samples), 200)}") if not per_position_shifts: return {'error': 'No predictions could be made'} shifts = [p['cls_shift'] for p in per_position_shifts] mean_shift = float(np.mean(shifts)) std_shift = float(np.std(shifts)) print(f" Total token replacements: {total_predictions}") print(f" Mean CLS shift: {mean_shift:.4f} ± {std_shift:.4f}") # Plot shift distribution plt = setup_nature_style() fig, ax = plt.subplots(figsize=(8, 5)) ax.hist(shifts, bins=50, color=NATURE_COLORS['blue'], alpha=0.8, edgecolor='white') ax.axvline(mean_shift, color=NATURE_COLORS['red'], linestyle='--', linewidth=1.5, label=f'Mean = {mean_shift:.3f}') ax.set_xlabel('CLS Embedding Shift (L2 norm)') ax.set_ylabel('Count') ax.set_title(f'Token Replacement → CLS Shift Distribution\n(n={total_predictions} replacements)') ax.legend() plt.tight_layout() plt.savefig(os.path.join(output_dir, 'mlm_zeroshot_probe.png'), dpi=300, bbox_inches='tight') plt.close() results = { 'total_replacements': total_predictions, 'mean_cls_shift': mean_shift, 'std_cls_shift': std_shift, 'median_cls_shift': float(np.median(shifts)), } with open(os.path.join(output_dir, 'mlm_zeroshot_probe.json'), 'w') as f: json.dump(results, f, indent=2, default=str) return results # ═══════════════════════════════════════════════════════════════════════ # PROBE 8: Token Importance (Leave-one-out CLS shift) # ═══════════════════════════════════════════════════════════════════════ def probe_token_importance(model, tokenizer, device, output_dir, max_samples=200, **kwargs): _cached_embs = kwargs.get("_cached_embs") _cached_samples = kwargs.get("_cached_samples") print("\n" + "="*60) print("PROBE 8: Token Importance (Leave-One-Out)") print("="*60) import torch import torch.nn.functional as F if _cached_samples is not None: samples = _cached_samples[:max_samples] else: samples = load_pretrain_wurcs(tokenizer, max_n=max_samples) tok = WURCSTokenizer(str(VOCAB_PATH)) MAX_LEN = 256 def _get_cls2(token_ids, branch_depths, linkage_types): """Helper: pad/truncate and run model.seq_embeddings(), return CLS numpy.""" min_l = min(len(token_ids), len(branch_depths), len(linkage_types)) token_ids = token_ids[:min_l] branch_depths = branch_depths[:min_l] linkage_types = linkage_types[:min_l] if min_l > MAX_LEN: token_ids = token_ids[:MAX_LEN] branch_depths = branch_depths[:MAX_LEN] linkage_types = linkage_types[:MAX_LEN] elif min_l < MAX_LEN: p = MAX_LEN - min_l token_ids = F.pad(token_ids, (0, p), value=0) branch_depths = F.pad(branch_depths, (0, p), value=0) linkage_types = F.pad(linkage_types, (0, p), value=0) with torch.no_grad(): out = model.seq_embeddings( token_ids.unsqueeze(0).to(device), branch_depths=branch_depths.unsqueeze(0).to(device), linkage_types=linkage_types.unsqueeze(0).to(device), ) return out[0, 0, :].cpu().numpy() # For each sample, drop one token at a time and measure CLS shift all_importance_by_position = {} # position -> list of shifts token_importance_map = {} # token_id -> list of shifts for idx, s in enumerate(samples): if idx > 100: break try: result = tok.tokenize(s['wurcs'], max_length=MAX_LEN) ids = torch.tensor(result['token_ids'], dtype=torch.long) bd = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long) lt = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long) real_len = result.get('length', len(ids)) if real_len < 4: continue # Original CLS cls_orig = _get_cls2(ids.clone(), bd.clone(), lt.clone()) seq_len = real_len for pos in range(1, min(seq_len - 1, 30)): # Remove token at position pos (drop from all 3 tensors) ids_dropped = torch.cat([ids[:pos], ids[pos+1:]]) bd_dropped = torch.cat([bd[:pos], bd[pos+1:]]) lt_dropped = torch.cat([lt[:pos], lt[pos+1:]]) cls_dropped = _get_cls2(ids_dropped, bd_dropped, lt_dropped) shift = float(np.linalg.norm(cls_orig - cls_dropped)) rel_pos = pos / seq_len # relative position # Track by relative position bin bin_key = f'{int(rel_pos * 10) / 10:.1f}' if bin_key not in all_importance_by_position: all_importance_by_position[bin_key] = [] all_importance_by_position[bin_key].append(shift) # Track by token id tid = ids[pos].item() if tid not in token_importance_map: token_importance_map[tid] = [] token_importance_map[tid].append(shift) except Exception as e: continue if idx % 25 == 0: print(f" Processed {idx}/100") if not all_importance_by_position: return {'error': 'No importance data'} # Average importance by position pos_importance = {k: float(np.mean(v)) for k, v in sorted(all_importance_by_position.items())} print(f" Position importance: {pos_importance}") # Top-10 most important tokens token_avg = {k: float(np.mean(v)) for k, v in token_importance_map.items() if len(v) >= 3} top_tokens = sorted(token_avg.items(), key=lambda x: x[1], reverse=True)[:20] print(f" Top-10 most important tokens (by CLS shift):") for tid, shift in top_tokens[:10]: print(f" Token {tid}: shift={shift:.4f}") # Plot: importance by position plt = setup_nature_style() fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5)) positions = sorted(pos_importance.keys()) imp_vals = [pos_importance[p] for p in positions] ax1.bar(range(len(positions)), imp_vals, color=NATURE_COLORS['blue'], alpha=0.8, edgecolor='white') ax1.set_xticks(range(len(positions))) ax1.set_xticklabels(positions, fontsize=8) ax1.set_xlabel('Relative Position in Sequence') ax1.set_ylabel('Mean CLS Shift') ax1.set_title('Token Importance by Position') # Plot: top token importance top_ids = [str(t[0]) for t in top_tokens[:15]] top_shifts = [t[1] for t in top_tokens[:15]] ax2.barh(range(len(top_ids)), top_shifts, color=NATURE_COLORS['orange'], alpha=0.8, edgecolor='white') ax2.set_yticks(range(len(top_ids))) ax2.set_yticklabels(top_ids, fontsize=8) ax2.set_xlabel('Mean CLS Shift') ax2.set_title('Top-15 Most Important Tokens') ax2.invert_yaxis() plt.suptitle('Token Importance Analysis (Leave-One-Out)', fontsize=13) plt.tight_layout() plt.savefig(os.path.join(output_dir, 'token_importance_probe.png'), dpi=300, bbox_inches='tight') plt.close() results = { 'position_importance': pos_importance, 'top_tokens': {str(k): v for k, v in top_tokens}, 'n_samples_processed': min(len(samples), 100), } with open(os.path.join(output_dir, 'token_importance_probe.json'), 'w') as f: json.dump(results, f, indent=2, default=str) return results # ═══════════════════════════════════════════════════════════════════════ # Main # ═══════════════════════════════════════════════════════════════════════ # ============================================================ # PROBE 9: Cancer Glycan Marker Signatures # ============================================================ def probe_cancer_markers(model, tokenizer, device, output_dir, max_samples=5000): # Create args-like namespace for compatibility import types args = types.SimpleNamespace(output_dir=output_dir, max_samples=max_samples, batch_size=64) """Probe whether embeddings separate cancer-associated glycan signatures. Cancer cells have aberrant glycosylation: hyper-sialylation (Neu5Ac), hyper-fucosylation (Fuc), and truncated O-glycans (Tn antigen = single GalNAc). We classify glycans as "cancer-associated" if they have >=2 sialylation markers OR specific truncation patterns. """ print("\n" + "="*60) print("PROBE 9: Cancer Glycan Marker Signatures") print("="*60) import csv, json from pathlib import Path root = Path(__file__).resolve().parent.parent.parent csv_path = root / 'bert_training_v4' / 'data' / 'multimodal_index.csv' samples = [] with open(csv_path) as fh: reader = csv.DictReader(fh) for i, row in enumerate(reader): if i >= args.max_samples: break w = row.get('wurcs', '') monos = row.get('monosaccharide_names', '') if not w: continue mono_list = [m.strip() for m in monos.split(',') if m.strip()] # Cancer-associated markers n_sialic = sum(1 for m in mono_list if m in ('Neu5Ac', 'Neu5Gc', 'KDN')) n_fuc = sum(1 for m in mono_list if m == 'Fuc') n_galnac = sum(1 for m in mono_list if m == 'GalNAc') total_monos = len(mono_list) # Cancer score: hyper-sialylation OR hyper-fucosylation OR truncated sialylation_ratio = n_sialic / max(total_monos, 1) fucosylation_ratio = n_fuc / max(total_monos, 1) is_truncated = (total_monos <= 2 and n_galnac >= 1) # Tn-like # Binary: cancer-associated if high sialylation/fucosylation or truncated cancer_assoc = (sialylation_ratio >= 0.3 or fucosylation_ratio >= 0.3 or is_truncated) label = 'cancer_associated' if cancer_assoc else 'normal' samples.append({'wurcs': w, 'label': label, 'n_sialic': n_sialic, 'n_fuc': n_fuc, 'sialylation_ratio': sialylation_ratio}) labels = [s['label'] for s in samples] from collections import Counter dist = Counter(labels) print(f" Total: {len(samples)}, Distribution: {dict(dist)}") if dist['cancer_associated'] < 20 or dist['normal'] < 20: print(" Too few samples in one class, skipping") return {} # Balance classes min_n = min(dist.values()) balanced = [] counts = {'cancer_associated': 0, 'normal': 0} for s in samples: if counts[s['label']] < min_n: balanced.append(s) counts[s['label']] += 1 print(f" Balanced: {len(balanced)} ({min_n} per class)") embs = batch_cls_embeddings(model, balanced, device=device, batch_size=args.batch_size if hasattr(args, 'batch_size') else 64) if embs is None or len(embs) == 0: print(" SKIPPING — no valid embeddings") return {} from sklearn.metrics import roc_auc_score, silhouette_score from sklearn.model_selection import cross_val_score from sklearn.linear_model import LogisticRegression X = np.array(embs) y = np.array([1 if s['label'] == 'cancer_associated' else 0 for s in balanced[:len(embs)]]) # Linear probe lr = LogisticRegression(max_iter=1000, random_state=42) scores = cross_val_score(lr, X, y, cv=5, scoring='roc_auc') mean_auc = scores.mean() std_auc = scores.std() # KNN from sklearn.neighbors import KNeighborsClassifier knn = KNeighborsClassifier(n_neighbors=5) knn_scores = cross_val_score(knn, X, y, cv=5, scoring='accuracy') # Silhouette try: sil = silhouette_score(X, y) except: sil = float('nan') results = { 'n_samples': len(embs), 'n_cancer': int(y.sum()), 'n_normal': int((1-y).sum()), 'linear_probe_auc': float(mean_auc), 'linear_probe_auc_std': float(std_auc), 'knn_accuracy': float(knn_scores.mean()), 'silhouette': float(sil), } print(f" Linear Probe AUC: {mean_auc:.4f} ± {std_auc:.4f}") print(f" KNN Accuracy: {knn_scores.mean():.4f}") print(f" Silhouette: {sil:.4f}") # Save results out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) with open(out_dir / 'cancer_markers_probe.json', 'w') as fh: json.dump(results, fh, indent=2) # Plot try: import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from sklearn.decomposition import PCA pca = PCA(n_components=2) X2 = pca.fit_transform(X) fig, ax = plt.subplots(1, 1, figsize=(8, 6)) for label_val, label_name, color in [(1, 'Cancer-associated', 'red'), (0, 'Normal', 'blue')]: mask = y == label_val ax.scatter(X2[mask, 0], X2[mask, 1], c=color, alpha=0.3, s=10, label=label_name) ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})') ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})') ax.set_title(f'Cancer Glycan Markers (AUC={mean_auc:.3f})') ax.legend() plt.tight_layout() plt.savefig(out_dir / 'cancer_markers_probe.png', dpi=150) plt.close() except Exception as e: print(f" Plot error: {e}") # Publication plots try: save_publication_plots(X, [s['label'] for s in balanced[:len(embs)]], 'Cancer Glycan Markers', out_dir, title_prefix='Probe 9') except Exception as e: print(f" Pub plot error: {e}") return results def probe_glycosylation_type(model, tokenizer, device, output_dir, max_samples=5000): # Create args-like namespace for compatibility import types args = types.SimpleNamespace(output_dir=output_dir, max_samples=max_samples, batch_size=64) """Probe whether embeddings separate N-linked vs O-linked vs free glycans. Uses curated GlycanML benchmark link data. """ print("\n" + "="*60) print("PROBE 10: Glycosylation Type (N vs O)") print("="*60) import csv, json from pathlib import Path root = Path(__file__).resolve().parent.parent.parent link_csv = root / 'bench' / 'GlycanML' / 'data' / 'glycan_link_wurcs_subset.csv' if not link_csv.exists(): print(f" Link data not found: {link_csv}") return {} samples = [] with open(link_csv) as fh: reader = csv.DictReader(fh) for row in reader: w = row.get('wurcs', '') link = row.get('link', '') if w and link in ('N', 'O'): samples.append({'wurcs': w, 'label': link}) from collections import Counter dist = Counter(s['label'] for s in samples) print(f" Samples: {len(samples)}, Distribution: {dict(dist)}") if len(samples) < 50: print(" Too few samples, skipping") return {} embs = batch_cls_embeddings(model, samples, device=device, batch_size=args.batch_size if hasattr(args, 'batch_size') else 64) if embs is None or len(embs) == 0: print(" SKIPPING — no valid embeddings") return {} from sklearn.metrics import silhouette_score from sklearn.neighbors import KNeighborsClassifier from sklearn.model_selection import cross_val_score from sklearn.preprocessing import LabelEncoder X = np.array(embs) le = LabelEncoder() y = le.fit_transform([s['label'] for s in samples[:len(embs)]]) classes = list(le.classes_) # KNN at multiple k results = {'n_samples': len(embs), 'distribution': dict(dist), 'classes': classes} for k in [5, 10, 20]: if len(embs) > k: knn = KNeighborsClassifier(n_neighbors=k) scores = cross_val_score(knn, X, y, cv=5, scoring='accuracy') results[f'knn_k{k}_accuracy'] = float(scores.mean()) print(f" KNN (k={k:2d}): {scores.mean():.4f}") # Silhouette try: sil = silhouette_score(X, y) results['silhouette'] = float(sil) print(f" Silhouette: {sil:.4f}") except: pass # Linear probe (N vs O only, binary) n_o_mask = np.array([s['label'] in ('N', 'O') for s in samples[:len(embs)]]) if n_o_mask.sum() > 50: from sklearn.linear_model import LogisticRegression X_no = X[n_o_mask] y_no = np.array([1 if s['label'] == 'N' else 0 for s in samples[:len(embs)]])[n_o_mask] lr = LogisticRegression(max_iter=1000, random_state=42) auc_scores = cross_val_score(lr, X_no, y_no, cv=5, scoring='roc_auc') results['n_vs_o_auc'] = float(auc_scores.mean()) results['n_vs_o_auc_std'] = float(auc_scores.std()) print(f" N-vs-O AUC: {auc_scores.mean():.4f} ± {auc_scores.std():.4f}") # Save out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) with open(out_dir / 'glycosylation_type_probe.json', 'w') as fh: json.dump(results, fh, indent=2) # Plot try: import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from sklearn.decomposition import PCA pca = PCA(n_components=2) X2 = pca.fit_transform(X) fig, ax = plt.subplots(1, 1, figsize=(8, 6)) colors = {'N': 'blue', 'O': 'red', 'free': 'green'} for c in classes: mask = np.array([s['label'] == c for s in samples[:len(embs)]]) ax.scatter(X2[mask, 0], X2[mask, 1], c=colors.get(c, 'gray'), alpha=0.4, s=15, label=f'{c} (n={mask.sum()})') ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})') ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})') ax.set_title('Glycosylation Type: N vs O') ax.legend() plt.tight_layout() plt.savefig(out_dir / 'glycosylation_type_probe.png', dpi=150) plt.close() except Exception as e: print(f" Plot error: {e}") # Publication plots try: save_publication_plots(X, [s['label'] for s in samples[:len(embs)]], 'Glycosylation Type', out_dir, title_prefix='Probe 10') except Exception as e: print(f" Pub plot error: {e}") return results def probe_taxonomic_class(model, tokenizer, device, output_dir, max_samples=5000): # Create args-like namespace for compatibility import types args = types.SimpleNamespace(output_dir=output_dir, max_samples=max_samples, batch_size=64) """Probe whether embeddings separate glycans by biological class. Uses GlycanML classification data with 90+ taxonomic classes. """ print("\n" + "="*60) print("PROBE 11: Taxonomic Classification (GlycanML)") print("="*60) import csv, json from pathlib import Path from collections import Counter root = Path(__file__).resolve().parent.parent.parent cls_csv = root / 'bench' / 'GlycanML' / 'data' / 'glycan_classification_wurcs_subset.csv' if not cls_csv.exists(): print(f" Classification data not found: {cls_csv}") return {} samples = [] with open(cls_csv) as fh: reader = csv.DictReader(fh) for row in reader: w = row.get('wurcs', '') cls_label = row.get('class', '').strip() domain = row.get('domain', '').strip() kingdom = row.get('kingdom', '').strip() phylum = row.get('phylum', '').strip() if w and cls_label: samples.append({ 'wurcs': w, 'class': cls_label, 'domain': domain, 'kingdom': kingdom, 'phylum': phylum }) # Filter to classes with >= 20 samples for meaningful evaluation class_dist = Counter(s['class'] for s in samples) valid_classes = {c for c, n in class_dist.items() if n >= 20} samples = [s for s in samples if s['class'] in valid_classes] class_dist = Counter(s['class'] for s in samples) print(f" Samples: {len(samples)}, Classes (n>=20): {len(valid_classes)}") print(f" Top-10: {class_dist.most_common(10)}") if len(samples) < 100 or len(valid_classes) < 3: print(" Too few samples/classes, skipping") return {} # Cap at max_samples if len(samples) > args.max_samples: samples = samples[:args.max_samples] embs = batch_cls_embeddings(model, samples, device=device, batch_size=args.batch_size if hasattr(args, 'batch_size') else 64) if embs is None or len(embs) == 0: print(" SKIPPING — no valid embeddings") return {} from sklearn.metrics import silhouette_score from sklearn.neighbors import KNeighborsClassifier from sklearn.model_selection import cross_val_score from sklearn.preprocessing import LabelEncoder X = np.array(embs) results = {'n_samples': len(embs), 'n_classes': len(valid_classes)} # Evaluate at class level le = LabelEncoder() y_class = le.fit_transform([s['class'] for s in samples[:len(embs)]]) for k in [5, 10, 20]: if len(embs) > k: knn = KNeighborsClassifier(n_neighbors=k) scores = cross_val_score(knn, X, y_class, cv=5, scoring='accuracy') results[f'class_knn_k{k}'] = float(scores.mean()) print(f" Class KNN (k={k:2d}): {scores.mean():.4f}") # Silhouette at class level try: sil = silhouette_score(X, y_class) results['class_silhouette'] = float(sil) print(f" Class Silhouette: {sil:.4f}") except: pass # Also evaluate at domain level (coarser, fewer classes) domain_labels = [s.get('domain', '') for s in samples[:len(embs)]] domain_dist = Counter(domain_labels) valid_domains = {d for d, n in domain_dist.items() if n >= 10 and d} if len(valid_domains) >= 2: domain_mask = np.array([s.get('domain', '') in valid_domains for s in samples[:len(embs)]]) le_d = LabelEncoder() y_domain = le_d.fit_transform([s.get('domain', '') for s in samples[:len(embs)] if s.get('domain', '') in valid_domains]) X_d = X[domain_mask] if len(X_d) > 20: knn_d = KNeighborsClassifier(n_neighbors=10) d_scores = cross_val_score(knn_d, X_d, y_domain, cv=5, scoring='accuracy') results['domain_knn_k10'] = float(d_scores.mean()) results['domain_distribution'] = dict(Counter( s.get('domain', '') for s in samples[:len(embs)] if s.get('domain', '') in valid_domains)) print(f" Domain KNN (k=10): {d_scores.mean():.4f}") # Save out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) with open(out_dir / 'taxonomic_class_probe.json', 'w') as fh: json.dump(results, fh, indent=2) # Plot — PCA colored by top-5 classes try: import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from sklearn.decomposition import PCA pca = PCA(n_components=2) X2 = pca.fit_transform(X) top5 = [c for c, _ in class_dist.most_common(5)] fig, ax = plt.subplots(1, 1, figsize=(10, 7)) colors_list = ['red', 'blue', 'green', 'orange', 'purple'] for idx, cls_name in enumerate(top5): mask = np.array([s['class'] == cls_name for s in samples[:len(embs)]]) ax.scatter(X2[mask, 0], X2[mask, 1], c=colors_list[idx], alpha=0.3, s=10, label=f'{cls_name} (n={mask.sum()})') # Plot rest in gray rest_mask = np.array([s['class'] not in top5 for s in samples[:len(embs)]]) ax.scatter(X2[rest_mask, 0], X2[rest_mask, 1], c='lightgray', alpha=0.1, s=5, label=f'Other ({rest_mask.sum()})') ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})') ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})') ax.set_title(f'Taxonomic Classification ({len(valid_classes)} classes)') ax.legend(fontsize=8) plt.tight_layout() plt.savefig(out_dir / 'taxonomic_class_probe.png', dpi=150) plt.close() except Exception as e: print(f" Plot error: {e}") # Publication plots — by domain (cleaner than 90 classes) try: domain_labels = [s.get('domain', 'unknown') for s in samples[:len(embs)]] save_publication_plots(X, domain_labels, 'Taxonomic Domain', out_dir, title_prefix='Probe 11a') # Also by top-5 classes from collections import Counter as Ctr top5cls = [c for c, _ in Ctr(s['class'] for s in samples[:len(embs)]).most_common(5)] labels_top5 = [s['class'] if s['class'] in top5cls else 'Other' for s in samples[:len(embs)]] save_publication_plots(X, labels_top5, 'Top-5 Taxonomic Classes', out_dir, title_prefix='Probe 11b') except Exception as e: print(f" Pub plot error: {e}") return results # ============================================================ # Probe Registry # ============================================================ PROBES = { # 'ambiguity': probe_ambiguity, # Deprecated: ? marks are annotation artifacts, not biology 'composition': probe_composition, 'knn_purity': probe_knn_purity, 'polymerization': probe_polymerization, 'size_prediction': probe_size, 'link_binary': probe_link_binary, 'mlm_zeroshot': probe_mlm_zeroshot, 'token_importance': probe_token_importance, 'cancer_markers': probe_cancer_markers, 'glycosylation_type': probe_glycosylation_type, 'taxonomic_class': probe_taxonomic_class, } def main(): parser = argparse.ArgumentParser() parser.add_argument('--model', choices=['v5', 'v6'], required=True) parser.add_argument('--probe', nargs='+', default=['all'], choices=['all'] + list(PROBES.keys())) parser.add_argument('--output_dir', default=None) parser.add_argument('--device', default='cuda') parser.add_argument('--resolved_only', action='store_true', help='Filter out ambiguous WURCS (containing ?) before probing') parser.add_argument('--max_samples', type=int, default=5000) args = parser.parse_args() if args.output_dir is None: args.output_dir = str(PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / f'novel_probes_v2_{args.model}') os.makedirs(args.output_dir, exist_ok=True) print(f"Loading tokenizer from {VOCAB_PATH}...") tokenizer = WURCSTokenizer(str(VOCAB_PATH)) print(f" Vocab size: {tokenizer.vocab_size}") ckpt = CHECKPOINTS[args.model] if not ckpt.exists(): print(f"ERROR: Checkpoint not found: {ckpt}") sys.exit(1) model = load_model(str(ckpt), device=args.device) # ============================================================ # PHASE 1: Embed pretraining samples ONCE, save to disk # ============================================================ cache_dir = Path(args.output_dir) / 'embedding_cache' cache_dir.mkdir(parents=True, exist_ok=True) emb_npy = cache_dir / 'cls_embeddings.npy' samples_pkl = cache_dir / 'samples.pkl' if emb_npy.exists() and samples_pkl.exists(): print(f"\n Loading cached embeddings from {cache_dir}...") pretrain_embs = np.load(str(emb_npy)) import pickle with open(samples_pkl, 'rb') as pf: pretrain_samples = pickle.load(pf) print(f" Loaded: {pretrain_embs.shape[0]} embeddings, {len(pretrain_samples)} samples") else: print(f"\n Embedding pretraining samples (embed once, save to disk)...") pretrain_samples = load_pretrain_wurcs(tokenizer, max_n=args.max_samples) if args.resolved_only: before = len(pretrain_samples) pretrain_samples = [s for s in pretrain_samples if '?' not in s.get('wurcs', '')] print(f" Filtered resolved: {before} -> {len(pretrain_samples)} (removed {before-len(pretrain_samples)} ambiguous)") print(f" Embedding {len(pretrain_samples)} samples...") pretrain_embs_list = batch_cls_embeddings(model, pretrain_samples, device=args.device) pretrain_embs = np.array(pretrain_embs_list) del pretrain_embs_list # Save to disk np.save(str(emb_npy), pretrain_embs) import pickle with open(samples_pkl, 'wb') as pf: pickle.dump(pretrain_samples, pf) print(f" Saved: {emb_npy} ({pretrain_embs.nbytes / 1e9:.2f} GB)") import gc gc.collect() print(f" Pretrain embeddings: shape={pretrain_embs.shape}") # ============================================================ # PHASE 2: Run probes # ============================================================ # Probes that use pretraining data get pre-computed embeddings # Probes that use external data (GlycanML) still call batch_cls_embeddings PRETRAIN_PROBES = {'composition', 'polymerization', 'size_prediction', 'mlm_zeroshot', 'token_importance', 'ambiguity'} EXTERNAL_PROBES = {'cancer_markers', 'glycosylation_type', 'taxonomic_class', 'knn_purity', 'link_binary'} probes_to_run = list(PROBES.keys()) if 'all' in args.probe else args.probe all_results = {} for pn in probes_to_run: try: if pn in PRETRAIN_PROBES: # Pass pre-computed embeddings — probe won't re-embed all_results[pn] = PROBES[pn]( model, tokenizer, args.device, args.output_dir, args.max_samples, _cached_embs=pretrain_embs, _cached_samples=pretrain_samples ) else: # External probes embed their own (small) datasets all_results[pn] = PROBES[pn]( model, tokenizer, args.device, args.output_dir, args.max_samples ) except Exception as e: print(f"\n ERROR in '{pn}': {e}") import traceback; traceback.print_exc() all_results[pn] = {'error': str(e)} gc.collect() # Free memory between probes # Save combined results with open(os.path.join(args.output_dir, 'all_probe_results_v2.json'), 'w') as f: json.dump(all_results, f, indent=2, default=str) print(f"\nALL PROBES COMPLETE — V2 ({args.model.upper()})") print(f"Results: {args.output_dir}") if __name__ == '__main__': main()