#!/usr/bin/env python3 """ Probe 7: Contrastive Learning Ablation — V5-A (MLM-only) vs V5.1 (MLM+SimCLR) Runs layer-wise CLS probing (Probe 2) and embedding vs structural similarity (Probe 4) on V5-A, then generates comparative overlay figures against V5.1. Usage: python probe_7_contrastive_ablation.py --device cuda python probe_7_contrastive_ablation.py --compare_only """ import os, sys, json, csv, argparse import numpy as np from pathlib import Path from collections import Counter PROJECT_ROOT = Path(__file__).resolve().parents[2] VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json' V5A_CKPT = PROJECT_ROOT / 'checkpoints_v5_bpe_topo' / 'best_v5_bpe_topo_model.pt' V51_CKPT = PROJECT_ROOT / 'bert_v5.1_contrastive' / 'checkpoints' / 'best_v51_contrastive_model.pt' BENCH_DIR = PROJECT_ROOT / 'bench' / 'GlycanML' / 'data' PROBE_DIR = PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / 'probe_results_v6' sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(PROJECT_ROOT / 'bert_training_v4')) sys.path.insert(0, str(PROJECT_ROOT / 'bert_v6_contrastive' / 'scripts')) from probe_2_layerwise_cls import ( load_model, extract_layerwise_cls, load_domain_data, load_glycosylation_data, load_immunogenicity_data, train_linear_probe, run_layerwise_probing ) from downstream_tasks.utils.tokenizer import WURCSTokenizer def extract_cls_final(model, samples, device='cuda', max_len=256): """Extract final-layer CLS embeddings.""" import torch, torch.nn.functional as F tokenizer = WURCSTokenizer(str(VOCAB_PATH)) embeddings = [] for si, s in enumerate(samples): try: result = tokenizer.tokenize(s['wurcs'], max_length=max_len) token_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) ml = min(len(token_ids), len(bd), len(lt)) token_ids, bd, lt = token_ids[:ml], bd[:ml], lt[:ml] if ml > max_len: token_ids, bd, lt = token_ids[:max_len], bd[:max_len], lt[:max_len] elif ml < max_len: pad = max_len - ml token_ids = F.pad(token_ids, (0, pad), value=0) bd = F.pad(bd, (0, pad), value=0) lt = F.pad(lt, (0, pad), value=0) with torch.no_grad(): hidden = model.seq_embeddings(token_ids.unsqueeze(0).to(device), branch_depths=bd.unsqueeze(0).to(device), linkage_types=lt.unsqueeze(0).to(device)) for layer in model.seq_layers: hidden = layer(hidden) embeddings.append(hidden[0, 0, :].cpu().numpy()) except: embeddings.append(np.zeros(768)) if si > 0 and si % 500 == 0: print(f" Processed {si}/{len(samples)}") return np.array(embeddings) def load_glyco_with_iupac(): csv_path = BENCH_DIR / 'glycan_link_wurcs_subset.csv' samples, labels, iupacs = [], [], [] with open(csv_path) as f: for row in csv.DictReader(f): w, link, iupac = row.get('wurcs',''), row.get('link',''), row.get('glycan','') if w.startswith('WURCS') and link in ('N','O') and iupac: samples.append({'wurcs': w}); labels.append(link); iupacs.append(iupac) print(f" Glyco with IUPAC: {len(samples)} samples") return samples, labels, iupacs def run_probe4(model, samples, iupacs, model_name, out_dir): from sklearn.metrics.pairwise import cosine_similarity print(f"\nExtracting CLS ({len(samples)} samples)...") embs = extract_cls_final(model, samples) embed_sim = cosine_similarity(embs) # Structural similarity cache = out_dir / 'struct_sim_cache.npy' if cache.exists(): print(" Loading cached struct sim...") struct_sim = np.load(str(cache)) else: try: from glycowork.motif.annotate import annotate_dataset print(f" Computing motif fingerprints...") mdf = annotate_dataset(iupacs, feature_set=['known']) struct_sim = cosine_similarity(mdf.values.astype(float)) np.save(str(cache), struct_sim) except Exception as e: print(f" ERROR: {e}"); return None n = len(samples) iu = np.triu_indices(n, k=1) from scipy.stats import pearsonr, spearmanr pr, pp = pearsonr(embed_sim[iu], struct_sim[iu]) sr, sp = spearmanr(embed_sim[iu], struct_sim[iu]) res = {'model': model_name, 'pearson_r': float(pr), 'spearman_rho': float(sr), 'n_pairs': int(len(embed_sim[iu])), 'n_glycans': n} print(f" {model_name}: Pearson r={pr:.4f}, Spearman rho={sr:.4f}") return res def plot_ablation(v5a_res, v51_res, out_dir): import matplotlib; matplotlib.use('Agg') import matplotlib.pyplot as plt plt.rcParams.update({'font.size': 11, 'figure.dpi': 300, 'axes.spines.top': False, 'axes.spines.right': False}) tasks = ['Domain', 'Immunogenicity', 'Glycosylation'] colors = {'Domain': '#E63946', 'Immunogenicity': '#457B9D', 'Glycosylation': '#2A9D8F'} fig, axes = plt.subplots(1, 3, figsize=(18, 5.5)) for i, task in enumerate(tasks): ax = axes[i] v5a_t = sorted([r for r in v5a_res if r['task']==task], key=lambda r: r['layer']) v51_t = sorted([r for r in v51_res if r['task']==task], key=lambda r: r['layer']) if not v5a_t or not v51_t: continue layers = [r['layer'] for r in v5a_t] v5a_a = [r['accuracy_mean'] for r in v5a_t] v51_a = [r['accuracy_mean'] for r in v51_t] v5a_s = [r['accuracy_std'] for r in v5a_t] v51_s = [r['accuracy_std'] for r in v51_t] ax.errorbar(layers, v5a_a, yerr=v5a_s, color='#888', marker='o', ms=6, lw=2, ls='--', capsize=3, label='V5-A (MLM only)', zorder=2) ax.fill_between(layers, [a-s for a,s in zip(v5a_a,v5a_s)], [a+s for a,s in zip(v5a_a,v5a_s)], color='#888', alpha=0.08) ax.errorbar(layers, v51_a, yerr=v51_s, color=colors[task], marker='s', ms=7, lw=2.5, capsize=3, label='V5.1 (MLM+SimCLR)', zorder=3) ax.fill_between(layers, [a-s for a,s in zip(v51_a,v51_s)], [a+s for a,s in zip(v51_a,v51_s)], color=colors[task], alpha=0.1) d = max(v51_a) - max(v5a_a) bl = v51_t[np.argmax(v51_a)]['layer'] ax.annotate(f'Delta = {"+" if d>=0 else ""}{d:.3f}', xy=(bl, max(v51_a)), xytext=(bl+1.5, max(v51_a)+0.01), fontsize=10, fontweight='bold', color=colors[task], arrowprops=dict(arrowstyle='->', color=colors[task])) ax.set_xlabel('Layer'); ax.set_ylabel('Accuracy') ax.set_title(task, fontweight='bold') ax.set_xticks(range(13)); ax.set_xticklabels(['Emb']+[str(j) for j in range(1,13)]) ax.legend(frameon=True, edgecolor='#ccc', loc='lower right') ax.grid(axis='y', alpha=0.3, ls='--') ax.text(-0.08, 1.05, f'({chr(97+i)})', transform=ax.transAxes, fontsize=14, fontweight='bold') fig.suptitle('Probe 7: Contrastive Pre-training Ablation', fontsize=15, fontweight='bold', y=1.02) plt.tight_layout() out = Path(out_dir) for fmt in ['png','pdf']: plt.savefig(out/f'probe7_ablation_layerwise.{fmt}', dpi=300, bbox_inches='tight', facecolor='white') print(f" Saved: {out/f'probe7_ablation_layerwise.{fmt}'}") plt.close() def plot_p4_comparison(v5a_p4, v51_p4, out_dir): import matplotlib; matplotlib.use('Agg') import matplotlib.pyplot as plt if not v5a_p4 or not v51_p4: print(" Skip P4 plot (missing data)"); return fig, ax = plt.subplots(1, 1, figsize=(6, 5)) models = ['V5-A\n(MLM only)', 'V5.1\n(MLM+SimCLR)'] rhos = [v5a_p4['spearman_rho'], v51_p4['spearman_rho']] rs = [v5a_p4['pearson_r'], v51_p4['pearson_r']] x = np.arange(2); w = 0.3 b1 = ax.bar(x-w/2, rhos, w, label='Spearman rho', color='#457B9D', zorder=3) b2 = ax.bar(x+w/2, rs, w, label='Pearson r', color='#A8DADC', zorder=3) for b in list(b1)+list(b2): ax.text(b.get_x()+b.get_width()/2, b.get_height()+0.005, f'{b.get_height():.3f}', ha='center', va='bottom', fontsize=11, fontweight='bold') dr = rhos[1]-rhos[0] ax.annotate(f'Delta_rho = {"+" if dr>=0 else ""}{dr:.3f}', xy=(1.5, max(rhos)+0.03), fontsize=12, fontweight='bold', color='#E63946', ha='center') ax.set_ylabel('Correlation'); ax.set_title('Embed-Structure Correlation', fontweight='bold') ax.set_xticks(x); ax.set_xticklabels(models, fontsize=12) ax.legend(frameon=True, edgecolor='#ccc'); ax.grid(axis='y', alpha=0.3, ls='--') ax.set_ylim(0, max(max(rhos), max(rs))+0.08) ax.spines['top'].set_visible(False); ax.spines['right'].set_visible(False) plt.tight_layout() out = Path(out_dir) for fmt in ['png','pdf']: plt.savefig(out/f'probe7_ablation_embed_vs_struct.{fmt}', dpi=300, bbox_inches='tight', facecolor='white') print(f" Saved: {out/f'probe7_ablation_embed_vs_struct.{fmt}'}") plt.close() def main(): parser = argparse.ArgumentParser() parser.add_argument('--device', default='cuda') parser.add_argument('--compare_only', action='store_true') parser.add_argument('--skip_probe4', action='store_true') args = parser.parse_args() out_dir = PROBE_DIR / 'probe_7_contrastive_ablation' out_dir.mkdir(parents=True, exist_ok=True) # Load existing V5.1 results v51_p2_path = PROBE_DIR / 'probe_2_layerwise_cls' / 'layerwise_results_v6.json' v51_p4_path = PROBE_DIR / 'probe_4_embed_vs_structure' / 'embed_vs_struct_v6.json' v51_results = json.load(open(v51_p2_path)) if v51_p2_path.exists() else [] v51_p4 = json.load(open(v51_p4_path)) if v51_p4_path.exists() else None print(f"V5.1 Probe 2: {len(v51_results)} entries") v5a_p2_path = out_dir / 'v5a_layerwise_results.json' v5a_p4_path = out_dir / 'v5a_probe4_results.json' if args.compare_only: v5a_results = json.load(open(v5a_p2_path)) if v5a_p2_path.exists() else [] v5a_p4 = json.load(open(v5a_p4_path)) if v5a_p4_path.exists() else None else: print(f"\n{'='*60}") print(f" PROBE 7a: Layer-wise CLS on V5-A (MLM-only)") print(f"{'='*60}") model = load_model(str(V5A_CKPT), device=args.device) print("\nLoading datasets...") dom_s, dom_l = load_domain_data() gly_s, gly_l = load_glycosylation_data() imm_s, imm_l = load_immunogenicity_data() if len(dom_s) > 3000: np.random.seed(42) idx = np.random.choice(len(dom_s), 3000, replace=False) dom_s = [dom_s[i] for i in idx]; dom_l = [dom_l[i] for i in idx] print(f" Subsampled domain to {len(dom_s)}") print("\nExtracting layer-wise CLS (V5-A)...") dom_emb = extract_layerwise_cls(model, dom_s, device=args.device) gly_emb = extract_layerwise_cls(model, gly_s, device=args.device) imm_emb = extract_layerwise_cls(model, imm_s, device=args.device) print("\nRunning linear probes...") v5a_results = (run_layerwise_probing(dom_emb, dom_l, 'Domain', 13) + run_layerwise_probing(imm_emb, imm_l, 'Immunogenicity', 13) + run_layerwise_probing(gly_emb, gly_l, 'Glycosylation', 13)) with open(v5a_p2_path, 'w') as f: json.dump(v5a_results, f, indent=2, default=str) csv_out = out_dir / 'v5a_layerwise_results.csv' with open(csv_out, 'w', newline='') as f: w = csv.DictWriter(f, fieldnames=['task','layer','accuracy_mean','accuracy_std','f1_mean','f1_std','n_samples','n_classes']) w.writeheader() for r in v5a_results: w.writerow({k: r[k] for k in w.fieldnames}) print(f" Saved: {v5a_p2_path}") # Probe 4 v5a_p4 = None if not args.skip_probe4: print(f"\n{'='*60}") print(f" PROBE 7b: Embed vs Structure on V5-A") print(f"{'='*60}") gs, _, iupacs = load_glyco_with_iupac() v5a_p4 = run_probe4(model, gs, iupacs, 'V5-A', out_dir) if v5a_p4: with open(v5a_p4_path, 'w') as f: json.dump(v5a_p4, f, indent=2) import gc, torch; del model; torch.cuda.empty_cache(); gc.collect() # Generate comparison figures print(f"\n{'='*60}") print(f" Generating Comparison Figures") print(f"{'='*60}") if v51_results and v5a_results: plot_ablation(v5a_results, v51_results, out_dir) if not v5a_p4 and v5a_p4_path.exists(): v5a_p4 = json.load(open(v5a_p4_path)) plot_p4_comparison(v5a_p4, v51_p4, out_dir) # Summary print(f"\n{'='*60}") print(f" ABLATION SUMMARY") print(f"{'='*60}") for task in ['Domain', 'Immunogenicity', 'Glycosylation']: v5a_t = [r for r in v5a_results if r['task']==task] v51_t = [r for r in v51_results if r['task']==task] if v5a_t and v51_t: v5a_b = max(v5a_t, key=lambda r: r['accuracy_mean']) v51_b = max(v51_t, key=lambda r: r['accuracy_mean']) d = v51_b['accuracy_mean'] - v5a_b['accuracy_mean'] print(f"\n{task}:") print(f" V5-A best: L{v5a_b['layer']} = {v5a_b['accuracy_mean']:.4f}") print(f" V5.1 best: L{v51_b['layer']} = {v51_b['accuracy_mean']:.4f}") print(f" Delta = {'+'if d>=0 else ''}{d:.4f} ({'+'if d>=0 else ''}{d*100:.2f}pp)") if v5a_p4 and v51_p4: dr = v51_p4['spearman_rho'] - v5a_p4['spearman_rho'] print(f"\nEmbed-Structure: V5-A rho={v5a_p4['spearman_rho']:.4f}, V5.1 rho={v51_p4['spearman_rho']:.4f}, Delta={'+'if dr>=0 else ''}{dr:.4f}") summary = {'v5a_probe2': v5a_results, 'v51_probe2': v51_results, 'v5a_probe4': v5a_p4, 'v51_probe4': v51_p4} with open(out_dir / 'ablation_summary.json', 'w') as f: json.dump(summary, f, indent=2, default=str) print(f"\nCOMPLETE") if __name__ == '__main__': main()