Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,10 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/envhog_phrog2/all_positive_viseq_pairs.csv filter=lfs diff=lfs merge=lfs -text
+data/envhog_phrog2/envhog_filtered_proteins.fasta filter=lfs diff=lfs merge=lfs -text
+data/envhog_phrog2/envhog_phrog2__low_thr_enriched_final.csv filter=lfs diff=lfs merge=lfs -text
+data/envhog_phrog2/envhog_test_filtered_proteins_annotated_with_no_nan.csv filter=lfs diff=lfs merge=lfs -text
+data/envhog_phrog2/envhog_test_final_no_leakage.csv filter=lfs diff=lfs merge=lfs -text
+data/envhog_phrog2/envhog_train_final_no_leakage.csv filter=lfs diff=lfs merge=lfs -text
+data/envhog_phrog2/envhog_train_final_no_leakage_no_nan.csv filter=lfs diff=lfs merge=lfs -text

data/envhog_phrog2/all_positive_viseq_pairs.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54a0de0f3904630947fb75b84f7d8133cca2dcdd831d04a526db979472ea94ca
+size 107596414

data/envhog_phrog2/envhog_filtered_proteins.fasta

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f3106d5aeec0a237466f270b7f5d6be237a60b0785ce20ca2f53dc87de84282d
+size 459033917

data/envhog_phrog2/envhog_phrog2__low_thr_enriched_final.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:813be5289fcc867ca29ea968110243b83909ae2e8cc651dba739c013cae51d8f
+size 23396438

data/envhog_phrog2/envhog_test_filtered_proteins_annotated_with_no_nan.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8317cb0e229544ecb46eb25540c784eb1fe9310146acddb2cf2640f090f96a17
+size 101224690

data/envhog_phrog2/envhog_test_final_no_leakage.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bbf2b9261ed883257449765eb3203f96b06307dba4bbb32ba3a028707519157
+size 182087311

data/envhog_phrog2/envhog_train_final_no_leakage.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a4805494ceb4499ff5a29d34fe66b61c42d762b3e210c3b1bfe1a7c1906081b
+size 95012046

data/envhog_phrog2/envhog_train_final_no_leakage_no_nan.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b3c03c6694c853e503b75ba32b4a50777bf4a002c0a2c353f5fc4510d4db752
+size 70167125

src/eval_EmbeddingSpace.py

@@ -0,0 +1,703 @@
+#!/usr/bin/env python3
+"""
+Embedding Space Evaluation — Base vs All Fine-tuned ProtT5 XL Models.
+
+Loads the base model + all fine-tuned variants (Here Just: ContraMLM) and produces exactly three publication-quality plots:
+
+  Plot 1 — 2×2 t-SNE grid, one panel per model, coloured by PhrogCat.
+  Plot 2 — 1×3 scatter of pairwise L2 distances (base x-axis, fine-tuned
+            y-axis) for the same ~N_PAIRS protein pairs across all subplots.
+  Plot 3 — Same as Plot 2 but using cosine similarity.
+"""
+
+import os
+import sys
+import json
+import glob
+import re
+import argparse
+from itertools import combinations
+from typing import Any, Dict, List
+
+import torch
+import pandas as pd
+import numpy as np
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+from sklearn.manifold import TSNE
+from sklearn.decomposition import PCA
+from sklearn.metrics import silhouette_score
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+from peft import PeftModel
+
+import warnings
+warnings.filterwarnings("ignore")
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+
+VERSIONS: List[str] = ["ContraMLM_v1_1"]
+ALL_MODEL_LABELS: List[str] = ["base"] + VERSIONS
+
+DEFAULT_BASE_MODEL_NAME = "Rostlab/prot_t5_xl_uniref50"
+DEFAULT_DATA_PATH       = "./data/envhog_phrog2/envhog_test_final_no_leakage.csv"
+DEFAULT_MAX_LENGTH      = 512
+DEFAULT_SAMPLE_SIZE     = 5000
+DEFAULT_BATCH_SIZE      = 2
+DEFAULT_RANDOM_STATE    = 42
+DEFAULT_N_PAIRS         = 150   # protein pairs used in scatter plots 2 & 3
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Multi-model embedding space evaluation for ProtT5 XL"
+    )
+    parser.add_argument("--base-model",      type=str, default=DEFAULT_BASE_MODEL_NAME)
+    parser.add_argument("--data-path",       type=str, default=DEFAULT_DATA_PATH)
+    parser.add_argument("--output-dir",      type=str, default="./evaluation_results_EmbeddingSpace")
+    parser.add_argument("--max-length",      type=int, default=DEFAULT_MAX_LENGTH)
+    parser.add_argument("--sample-size",     type=int, default=DEFAULT_SAMPLE_SIZE)
+    parser.add_argument("--batch-size",      type=int, default=DEFAULT_BATCH_SIZE)
+    parser.add_argument("--random-state",    type=int, default=DEFAULT_RANDOM_STATE)
+    parser.add_argument(
+        "--n-pairs",
+        type=int,
+        default=DEFAULT_N_PAIRS,
+        help="Number of protein pairs for L2/cosine scatter plots (100–200 recommended)",
+    )
+    return parser.parse_args()
+
+
+args = parse_args()
+
+BASE_MODEL_NAME = args.base_model
+DATA_PATH       = args.data_path
+OUTPUT_DIR      = args.output_dir
+IMAGES_DIR      = os.path.join(OUTPUT_DIR, "images")
+TEXT_DIR        = os.path.join(OUTPUT_DIR, "text")
+MAX_LENGTH      = args.max_length
+SAMPLE_SIZE     = args.sample_size
+BATCH_SIZE      = args.batch_size
+RANDOM_STATE    = args.random_state
+N_PAIRS_TARGET  = max(2, args.n_pairs)
+
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+os.makedirs(IMAGES_DIR, exist_ok=True)
+os.makedirs(TEXT_DIR,   exist_ok=True)
+
+
+# ============================================================================
+# LOGGING
+# ============================================================================
+
+class Tee:
+    def __init__(self, *streams):
+        self.streams = streams
+    def write(self, data):
+        for s in self.streams:
+            s.write(data)
+            s.flush()
+    def flush(self):
+        for s in self.streams:
+            s.flush()
+
+
+run_log_path = os.path.join(TEXT_DIR, "run_log.txt")
+log_file     = open(run_log_path, "w", encoding="utf-8")
+sys.stdout   = Tee(sys.__stdout__, log_file)
+sys.stderr   = Tee(sys.__stderr__, log_file)
+
+print("=" * 80)
+print("MULTI-MODEL EMBEDDING SPACE EVALUATION")
+print("=" * 80)
+print(f"\nBase model  : {BASE_MODEL_NAME}")
+print(f"Versions    : {VERSIONS}")
+print(f"Data        : {DATA_PATH}")
+print(f"Sample size : {SAMPLE_SIZE}")
+print(f"N pairs     : {N_PAIRS_TARGET}")
+print(f"Output      : {OUTPUT_DIR}")
+
+# ============================================================================
+# DEVICE
+# ============================================================================
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"\nDevice: {device}")
+if torch.cuda.is_available():
+    print(f"GPU   : {torch.cuda.get_device_name(0)}")
+
+# ============================================================================
+# LOAD DATA
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("LOADING DATA")
+print("=" * 80)
+
+df = pd.read_csv(DATA_PATH)
+print(f"Total proteins: {len(df)}")
+
+if SAMPLE_SIZE > 0 and len(df) > SAMPLE_SIZE:
+    df_sample = df.sample(n=SAMPLE_SIZE, random_state=RANDOM_STATE).reset_index(drop=True)
+    print(f"Sampled {SAMPLE_SIZE} proteins")
+else:
+    df_sample = df.reset_index(drop=True)
+    print(f"Using all {len(df)} proteins")
+
+sequences   = df_sample["sequence"].tolist()
+lengths     = df_sample["length"].tolist()
+protein_ids = df_sample["id"].tolist()
+
+# PhrogCat — used as colour label in t-SNE
+if "PhrogCat" in df_sample.columns:
+    phrog_cats = df_sample["PhrogCat"].fillna("unknown").tolist()
+else:
+    print("WARNING: 'PhrogCat' column not found — using 'unknown' for all proteins")
+    phrog_cats = ["unknown"] * len(df_sample)
+
+
+def prepare_t5_seq(seq: str) -> str:
+    return " ".join(list(str(seq).replace(" ", "")))
+
+
+sequences = [prepare_t5_seq(s) for s in sequences]
+
+# ============================================================================
+# TOKENIZER
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("LOADING TOKENIZER")
+print("=" * 80)
+
+tokenizer = T5Tokenizer.from_pretrained(
+    BASE_MODEL_NAME, do_lower_case=False, legacy=True
+)
+print(f"Tokenizer loaded: {BASE_MODEL_NAME}")
+
+# ============================================================================
+# HELPERS
+# ============================================================================
+
+def get_encoder(model):
+    """Return the encoder module regardless of wrapper type."""
+    if hasattr(model, "encoder"):
+        return model.encoder
+    fn = getattr(model, "get_encoder", None)
+    if callable(fn):
+        return fn()
+    for attr in ("base_model", "model"):
+        inner = getattr(model, attr, None)
+        if inner is not None:
+            if hasattr(inner, "encoder"):
+                return inner.encoder
+            fn2 = getattr(inner, "get_encoder", None)
+            if callable(fn2):
+                return fn2()
+    return model
+
+
+def select_best_adapter_dir(version: str) -> str:
+    """Return the first valid (NaN/Inf-free) adapter directory for *version*."""
+    finetuned_path  = f"./runs/protrans_XL_Full_lora_envhog_{version}/lora_adapters"
+    checkpoint_root = f"./runs/protrans_XL_Full_lora_envhog_{version}"
+
+    candidates = []
+    if os.path.isdir(finetuned_path):
+        candidates.append(finetuned_path)
+
+    ckpt_paths = sorted(
+        glob.glob(os.path.join(checkpoint_root, "checkpoint-*")),
+        key=lambda p: int(re.search(r"checkpoint-(\d+)", p).group(1))
+        if re.search(r"checkpoint-(\d+)", p) else -1,
+        reverse=True,
+    )
+    for cp in ckpt_paths:
+        if os.path.isdir(cp):
+            candidates.append(cp)
+
+    def _resolve(candidate):
+        for subdir in (candidate, os.path.join(candidate, "lora_adapters")):
+            if (os.path.isfile(os.path.join(subdir, "adapter_model.safetensors"))
+                    or os.path.isfile(os.path.join(subdir, "adapter_model.bin"))):
+                return subdir
+        return None
+
+    for candidate in candidates:
+        resolved = _resolve(candidate)
+        if resolved is None:
+            continue
+        safe = os.path.join(resolved, "adapter_model.safetensors")
+        bin_ = os.path.join(resolved, "adapter_model.bin")
+        if os.path.isfile(safe):
+            from safetensors.torch import load_file
+            state_dict = load_file(safe, device="cpu")
+        else:
+            state_dict = torch.load(bin_, map_location="cpu")
+
+        has_nan = any(torch.isnan(v).any().item() for v in state_dict.values())
+        has_inf = any(torch.isinf(v).any().item() for v in state_dict.values())
+        if not has_nan and not has_inf:
+            print(f"  [{version}] Using adapter: {resolved}")
+            return resolved
+
+    raise RuntimeError(
+        f"No valid (NaN/Inf-free) adapter found for version '{version}'. "
+        f"Searched: {candidates}"
+    )
+
+
+def get_embeddings(model, seqs: List[str], batch_size: int = 16) -> np.ndarray:
+    """Mean-pooled encoder embeddings for a list of pre-formatted sequences."""
+    encoder = get_encoder(model)
+    encoder.eval()
+    all_embs = []
+
+    with torch.no_grad():
+        for i in range(0, len(seqs), batch_size):
+            batch = seqs[i : i + batch_size]
+            inputs = tokenizer(
+                batch,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=MAX_LENGTH,
+            )
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            out    = encoder(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+            )
+            hidden = out.last_hidden_state
+            mask   = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
+            pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
+            all_embs.append(pooled.float().cpu().numpy())
+
+            if (i // batch_size) % 10 == 0:
+                print(f"  {i}/{len(seqs)} sequences processed...")
+
+    return np.vstack(all_embs)
+
+
+# ============================================================================
+# LOAD BASE MODEL
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("LOADING BASE MODEL")
+print("=" * 80)
+
+model_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+base_model = T5ForConditionalGeneration.from_pretrained(
+    BASE_MODEL_NAME, torch_dtype=model_dtype, low_cpu_mem_usage=True
+)
+base_model = base_model.to(device)
+base_model.eval()
+print("Base model loaded")
+
+# ============================================================================
+# GENERATE BASE EMBEDDINGS
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("GENERATING BASE EMBEDDINGS")
+print("=" * 80)
+
+base_embeddings = get_embeddings(base_model, sequences, BATCH_SIZE)
+print(f"Base embeddings shape: {base_embeddings.shape}")
+
+# Free base model GPU memory before loading fine-tuned models one-by-one
+# (keep the numpy array — it is small)
+del base_model
+torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+# ============================================================================
+# LOAD FINE-TUNED MODELS AND GENERATE EMBEDDINGS
+# ============================================================================
+
+# Stores: { version_name: np.ndarray }
+ft_embeddings: Dict[str, np.ndarray] = {}
+
+for version in VERSIONS:
+    print("\n" + "=" * 80)
+    print(f"LOADING FINE-TUNED MODEL: {version}")
+    print("=" * 80)
+
+    adapter_dir = select_best_adapter_dir(version)
+
+    ft_model = T5ForConditionalGeneration.from_pretrained(
+        BASE_MODEL_NAME, torch_dtype=model_dtype, low_cpu_mem_usage=True
+    )
+    ft_model = PeftModel.from_pretrained(ft_model, adapter_dir)
+
+    merge_fn = getattr(ft_model, "merge_and_unload", None)
+    if callable(merge_fn):
+        ft_model = merge_fn()
+        print("  LoRA adapters merged")
+
+    ft_model = ft_model.to(device)
+    ft_model.eval()
+
+    print(f"  Generating embeddings for {version}...")
+    ft_embeddings[version] = get_embeddings(ft_model, sequences, BATCH_SIZE)
+    print(f"  {version} embeddings shape: {ft_embeddings[version].shape}")
+
+    del ft_model
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
+# ============================================================================
+# SHARED t-SNE COLOUR MAP (PhrogCat)
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("PREPARING t-SNE COLOUR MAP")
+print("=" * 80)
+
+unique_cats   = sorted(set(phrog_cats))
+n_cats        = len(unique_cats)
+cmap_name     = "tab20" if n_cats > 10 else "tab10"
+cmap          = plt.get_cmap(cmap_name, n_cats)
+cat_to_idx    = {cat: i for i, cat in enumerate(unique_cats)}
+colour_values = np.array([cat_to_idx[c] for c in phrog_cats])
+
+print(f"Unique PhrogCat categories: {n_cats}")
+print(f"Categories: {unique_cats}")
+
+# ============================================================================
+# t-SNE FOR ALL FOUR MODELS
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("RUNNING t-SNE (4 MODELS)")
+print("=" * 80)
+
+n_tsne     = min(2000, len(sequences))
+rng_tsne   = np.random.default_rng(RANDOM_STATE)
+tsne_idx   = rng_tsne.choice(len(sequences), n_tsne, replace=False)
+perplexity = 30 if n_tsne > 30 else max(5, n_tsne - 1)
+
+# Collect all embedding matrices for the 4 models
+all_embeddings_ordered: Dict[str, np.ndarray] = {
+    "base": base_embeddings,
+    **ft_embeddings,
+}
+
+tsne_results: Dict[str, np.ndarray] = {}
+tsne_silhouette: Dict[str, float] = {}
+
+for label, emb in all_embeddings_ordered.items():
+    print(f"  PCA → t-SNE for [{label}]...")
+    pca   = PCA(n_components=50, random_state=RANDOM_STATE)
+    emb_pca = pca.fit_transform(emb)
+    tsne  = TSNE(n_components=2, random_state=RANDOM_STATE, perplexity=perplexity)
+    tsne_results[label] = tsne.fit_transform(emb_pca[tsne_idx])
+
+    # Silhouette index on the 2D t-SNE map using PhrogCat categories as labels.
+    labels_tsne = np.array(phrog_cats, dtype=object)[tsne_idx]
+    n_label_values = len(set(labels_tsne.tolist()))
+    if 2 <= n_label_values < len(labels_tsne):
+        try:
+            tsne_silhouette[label] = float(silhouette_score(tsne_results[label], labels_tsne))
+        except Exception:
+            tsne_silhouette[label] = float("nan")
+    else:
+        tsne_silhouette[label] = float("nan")
+
+    print(f"    Done.")
+    if np.isnan(tsne_silhouette[label]):
+        print("    Silhouette(PhrogCat): n/a")
+    else:
+        print(f"    Silhouette(PhrogCat): {tsne_silhouette[label]:.4f}")
+
+colours_tsne = colour_values[tsne_idx]
+
+# ============================================================================
+# SHARED PROTEIN PAIRS FOR SCATTER PLOTS
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("BUILDING SHARED PROTEIN PAIRS")
+print("=" * 80)
+
+# We need N_PAIRS_TARGET pairs from a pool of proteins.
+# Minimum proteins needed so combinations >= N_PAIRS_TARGET:
+#   n*(n-1)/2 >= N_PAIRS_TARGET  →  n ≈ ceil((1 + sqrt(1+8k))/2)
+import math
+n_prot_needed = math.ceil((1 + math.sqrt(1 + 8 * N_PAIRS_TARGET)) / 2)
+n_prot_needed = max(n_prot_needed, 2)
+n_prot_needed = min(n_prot_needed, len(sequences))
+
+rng_pairs     = np.random.default_rng(RANDOM_STATE + 1)
+pair_indices  = rng_pairs.choice(len(sequences), n_prot_needed, replace=False)
+pair_indices  = pair_indices.tolist()
+
+all_pairs     = list(combinations(pair_indices, 2))
+
+# Randomly subsample to exactly N_PAIRS_TARGET pairs if we have more
+if len(all_pairs) > N_PAIRS_TARGET:
+    rng_sub   = np.random.default_rng(RANDOM_STATE + 2)
+    chosen    = rng_sub.choice(len(all_pairs), N_PAIRS_TARGET, replace=False)
+    all_pairs = [all_pairs[i] for i in chosen]
+
+n_pairs_actual = len(all_pairs)
+print(f"Protein pool size : {n_prot_needed}")
+print(f"Pairs generated   : {n_pairs_actual}")
+
+
+def pairwise_l2(emb: np.ndarray, pairs: list) -> np.ndarray:
+    return np.array([
+        np.linalg.norm(emb[a] - emb[b])
+        for a, b in pairs
+    ])
+
+
+def pairwise_cosine(emb: np.ndarray, pairs: list, eps: float = 1e-12) -> np.ndarray:
+    sims = []
+    for a, b in pairs:
+        va, vb = emb[a], emb[b]
+        denom  = np.linalg.norm(va) * np.linalg.norm(vb)
+        sim    = np.dot(va, vb) / max(denom, eps)
+        sims.append(float(np.clip(sim, -1.0, 1.0)))
+    return np.array(sims)
+
+
+# Compute for base
+base_pair_l2  = pairwise_l2(base_embeddings,  all_pairs)
+base_pair_cos = pairwise_cosine(base_embeddings, all_pairs)
+
+# Compute for each fine-tuned version
+ft_pair_l2:  Dict[str, np.ndarray] = {}
+ft_pair_cos: Dict[str, np.ndarray] = {}
+
+for version in VERSIONS:
+    emb = ft_embeddings[version]
+    ft_pair_l2[version]  = pairwise_l2(emb,  all_pairs)
+    ft_pair_cos[version] = pairwise_cosine(emb, all_pairs)
+
+# ============================================================================
+# PLOT 1 — 2×2 t-SNE GRID  (coloured by PhrogCat)
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("PLOT 1: 2×2 t-SNE GRID")
+print("=" * 80)
+
+fig, axes = plt.subplots(2, 2, figsize=(16, 14))
+axes_flat  = axes.flatten()
+
+panel_order = ["base", "ContraMLM_v1_0", "Default_v2_1", "MLP_v0"]
+panel_titles = {
+    "base":           "Base Model",
+    "ContraMLM_v1_0": "ContraMLM v1.0",
+    "Default_v2_1":   "Default v2.1",
+    "MLP_v0":         "MLP v0",
+}
+
+for ax, label in zip(axes_flat, panel_order):
+    xy   = tsne_results[label]
+    sc   = ax.scatter(
+        xy[:, 0], xy[:, 1],
+        c=colours_tsne,
+        cmap=cmap_name,
+        vmin=0, vmax=n_cats - 1,
+        alpha=0.65,
+        s=8,
+        linewidths=0,
+    )
+    sil_txt = (
+        f"Silhouette(PhrogCat): {tsne_silhouette[label]:.3f}"
+        if not np.isnan(tsne_silhouette[label])
+        else "Silhouette(PhrogCat): n/a"
+    )
+    ax.set_title(
+        f"{panel_titles[label]}\n{sil_txt}",
+        fontsize=14,
+        fontweight="bold",
+        pad=8,
+    )
+    ax.set_xlabel("t-SNE 1", fontsize=10)
+    ax.set_ylabel("t-SNE 2", fontsize=10)
+    ax.tick_params(labelsize=8)
+
+# Shared legend for PhrogCat categories
+legend_handles = [
+    mpatches.Patch(color=cmap(cat_to_idx[cat] / max(n_cats - 1, 1)), label=cat)
+    for cat in unique_cats
+]
+fig.legend(
+    handles=legend_handles,
+    title="PhrogCat",
+    title_fontsize=14,
+    fontsize=12,
+    loc="lower center",
+    ncol=min(n_cats, 6),
+    bbox_to_anchor=(0.5, -0.02),
+    frameon=True,
+)
+
+fig.suptitle(
+    f"t-SNE Embedding Space — Base vs Fine-tuned Models\n"
+    f"(n={n_tsne} proteins, coloured by PhrogCat)",
+    fontsize=15,
+    fontweight="bold",
+    y=1.01,
+)
+plt.tight_layout()
+
+plot1_path = os.path.join(IMAGES_DIR, "plot1_tsne_4models.png")
+fig.savefig(plot1_path, dpi=300, bbox_inches="tight")
+plt.close(fig)
+print(f"Saved: {plot1_path}")
+
+# ============================================================================
+# PLOT 2 — PAIRWISE L2 SCATTER (base x-axis, fine-tuned y-axis)
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("PLOT 2: PAIRWISE L2 DISTANCE SCATTER")
+print("=" * 80)
+
+fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+
+version_titles = {
+    "ContraMLM_v1_0": "ContraMLM v1.0",
+    "Default_v2_1":   "Default v2.1",
+    "MLP_v0":         "MLP v0",
+}
+
+for ax, version in zip(axes, VERSIONS):
+    x = base_pair_l2
+    y = ft_pair_l2[version]
+
+    # Diagonal reference line
+    lim_min = min(x.min(), y.min()) * 0.98
+    lim_max = max(x.max(), y.max()) * 1.02
+    ax.plot([lim_min, lim_max], [lim_min, lim_max],
+            color="gray", linestyle="--", linewidth=1.0, alpha=0.7, label="y = x")
+
+    ax.scatter(x, y, alpha=0.55, s=20, color="#2E86AB", linewidths=0)
+
+    # Pearson r annotation
+    r = float(np.corrcoef(x, y)[0, 1])
+    ax.text(
+        0.05, 0.93, f"r = {r:.3f}",
+        transform=ax.transAxes,
+        fontsize=10,
+        verticalalignment="top",
+        bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
+    )
+
+    ax.set_xlim(lim_min, lim_max)
+    ax.set_ylim(lim_min, lim_max)
+    ax.set_xlabel("Base model — L2 distance", fontsize=11)
+    ax.set_ylabel(f"{version_titles[version]} — L2 distance", fontsize=11)
+    ax.set_title(f"L2: Base vs {version_titles[version]}", fontsize=13, fontweight="bold")
+    ax.set_aspect("equal", adjustable="box")
+    ax.grid(True, alpha=0.25)
+    ax.legend(fontsize=9)
+
+fig.suptitle(
+    f"Pairwise L2 Distance: Base vs Fine-tuned Models\n"
+    f"({n_pairs_actual} protein pairs, same pairs across all subplots)",
+    fontsize=14,
+    fontweight="bold",
+)
+plt.tight_layout()
+
+plot2_path = os.path.join(IMAGES_DIR, "plot2_pairwise_l2_scatter.png")
+fig.savefig(plot2_path, dpi=300, bbox_inches="tight")
+plt.close(fig)
+print(f"Saved: {plot2_path}")
+
+# ============================================================================
+# PLOT 3 — PAIRWISE COSINE SIMILARITY SCATTER
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("PLOT 3: PAIRWISE COSINE SIMILARITY SCATTER")
+print("=" * 80)
+
+fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+
+for ax, version in zip(axes, VERSIONS):
+    x = base_pair_cos
+    y = ft_pair_cos[version]
+
+    lim_min = min(x.min(), y.min()) - 0.02
+    lim_max = max(x.max(), y.max()) + 0.02
+    ax.plot([lim_min, lim_max], [lim_min, lim_max],
+            color="gray", linestyle="--", linewidth=1.0, alpha=0.7, label="y = x")
+
+    ax.scatter(x, y, alpha=0.55, s=20, color="#E84855", linewidths=0)
+
+    r = float(np.corrcoef(x, y)[0, 1])
+    ax.text(
+        0.05, 0.93, f"r = {r:.3f}",
+        transform=ax.transAxes,
+        fontsize=10,
+        verticalalignment="top",
+        bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
+    )
+
+    ax.set_xlim(lim_min, lim_max)
+    ax.set_ylim(lim_min, lim_max)
+    ax.set_xlabel("Base model — cosine similarity", fontsize=11)
+    ax.set_ylabel(f"{version_titles[version]} — cosine similarity", fontsize=11)
+    ax.set_title(f"Cosine: Base vs {version_titles[version]}", fontsize=13, fontweight="bold")
+    ax.set_aspect("equal", adjustable="box")
+    ax.grid(True, alpha=0.25)
+    ax.legend(fontsize=9)
+
+fig.suptitle(
+    f"Pairwise Cosine Similarity: Base vs Fine-tuned Models\n"
+    f"({n_pairs_actual} protein pairs, same pairs across all subplots)",
+    fontsize=14,
+    fontweight="bold",
+)
+plt.tight_layout()
+
+plot3_path = os.path.join(IMAGES_DIR, "plot3_pairwise_cosine_scatter.png")
+fig.savefig(plot3_path, dpi=300, bbox_inches="tight")
+plt.close(fig)
+print(f"Saved: {plot3_path}")
+
+# ============================================================================
+# SAVE PAIR DATA AS CSV (reproducibility)
+# ============================================================================
+
+pair_records = []
+for k, (a, b) in enumerate(all_pairs):
+    row = {
+        "pair_index":    k,
+        "protein_id_a":  protein_ids[a],
+        "protein_id_b":  protein_ids[b],
+        "base_l2":       float(base_pair_l2[k]),
+        "base_cosine":   float(base_pair_cos[k]),
+    }
+    for version in VERSIONS:
+        row[f"{version}_l2"]     = float(ft_pair_l2[version][k])
+        row[f"{version}_cosine"] = float(ft_pair_cos[version][k])
+    pair_records.append(row)
+
+pairs_csv_path = os.path.join(TEXT_DIR, "pairwise_distances_all_models.csv")
+pd.DataFrame(pair_records).to_csv(pairs_csv_path, index=False)
+print(f"\nPair data saved: {pairs_csv_path}")
+
+# ============================================================================
+# DONE
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("EVALUATION COMPLETE")
+print("=" * 80)
+print(f"\nOutputs written to : {OUTPUT_DIR}")
+print(f"  Plot 1 (t-SNE)   : {plot1_path}")
+print(f"  Plot 2 (L2)      : {plot2_path}")
+print(f"  Plot 3 (cosine)  : {plot3_path}")
+print(f"  Pair CSV         : {pairs_csv_path}")
+print(f"  Run log          : {run_log_path}")

src/train.py

@@ -0,0 +1,1010 @@
+import os
+import re
+import sys
+from collections import defaultdict
+from typing import cast
+import pandas as pd
+import numpy as np
+import inspect
+import random
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import (
+    T5Tokenizer,
+    T5ForConditionalGeneration,
+    TrainingArguments,
+    Trainer,
+    TrainerCallback,
+)
+from transformers.optimization import Adafactor, AdafactorSchedule
+from peft import get_peft_model, LoraConfig, TaskType
+import matplotlib.pyplot as plt
+
+
+print("=" * 80)
+print("PROTRANS LORA FINE-TUNING: CONTRASTIVE + MLM LOSS (ContraMLM v1)")
+print("=" * 80)
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("CONFIGURATION")
+print("=" * 80)
+
+# Model configuration
+MODEL_NAME = "Rostlab/prot_t5_xl_uniref50"
+
+# LoRA configuration
+LORA_R = 32
+LORA_ALPHA = 64
+LORA_DROPOUT = 0.1
+LORA_TARGET_MODULES = ["q", "k", "v", "o"]
+LORA_TASK_TYPE = TaskType.SEQ_2_SEQ_LM
+
+# Training configuration
+#   BATCH_SIZE = number of proteins drawn from the dataset per forward pass.
+#   For each non-orphan, the collator samples 1 positive on-the-fly, so the
+#   actual forward-pass batch has BATCH_SIZE..2*BATCH_SIZE unique proteins.
+BATCH_SIZE = 32
+GRADIENT_ACCUMULATION_STEPS = 2
+NUM_EPOCHS = 2
+MAX_LENGTH = 512
+NOISE_DENSITY = 0.15
+
+# Contrastive loss — Contrastive (Con) with full adjacency matrix
+#   Total loss = (1 - CONTRASTIVE_LAMBDA) * MLM_loss + CONTRASTIVE_LAMBDA * Con_loss
+#   A per-batch adjacency matrix (N×N) is built from the VISEQ pair graph.
+#   Every known positive pair in the batch contributes to the numerator;
+#   false negatives are impossible by construction (adj built from ground-truth graph).
+#   CONTRASTIVE_LAMBDA is the convex-combination weight for the contrastive term.
+CONTRASTIVE_LAMBDA = 0.2
+CONTRASTIVE_TEMPERATURE = 0.1   # lower → sharper distribution → harder loss
+
+# Curriculum settings (same flags as Default_v2, applied to MLM component only)
+USE_LOSS_CLIPPING_CURRICULUM = False
+NUM_STAGES = 10
+KEEP_FRACTION_START = 0.20
+KEEP_FRACTION_END = 1.00
+LARGEST = False  # False = keep easiest losses first
+
+version = "v1_1"
+
+OUTPUT_DIR = f"./runs/protrans_XL_Full_lora_envhog_ContraMLM_{version}"
+
+# Data paths
+FASTA_FILE = "./data/envhog_phrog2/envhog_filtered_proteins.fasta"
+CSV_FILE   = "./data/envhog_phrog2/envhog_phrog2__low_thr_enriched_final.csv"
+PAIRS_FILE = "./data/envhog_phrog2/all_positive_viseq_pairs.csv"
+
+# Sample caps (applied to the pair list; existing ~168 K pairs are well below cap)
+MAX_TRAIN_SAMPLES = 400000
+MAX_EVAL_SAMPLES  = 50000
+
+print(f"Model:                    {MODEL_NAME}")
+print(f"LoRA rank:                {LORA_R}")
+print(f"Batch size (drawn):       {BATCH_SIZE}  ({BATCH_SIZE}–{2*BATCH_SIZE} unique proteins per step)")
+print(f"Gradient accumulation:    {GRADIENT_ACCUMULATION_STEPS}")
+print(f"Epochs:                   {NUM_EPOCHS}")
+print(f"Max sequence length:      {MAX_LENGTH}")
+print(f"Noise density (MLM):      {NOISE_DENSITY}")
+print(f"Contrastive lambda:       {CONTRASTIVE_LAMBDA}")
+print(f"Contrastive temperature:  {CONTRASTIVE_TEMPERATURE}")
+print(f"Loss clipping curriculum: {USE_LOSS_CLIPPING_CURRICULUM}")
+print(f"Output directory:         {OUTPUT_DIR}")
+
+# ============================================================================
+# CHECK PYTORCH AND GPU
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("GPU STATUS")
+print("=" * 80)
+
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available:  {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"CUDA device:     {torch.cuda.get_device_name(0)}")
+    print(f"Number of GPUs:  {torch.cuda.device_count()}")
+
+# ============================================================================
+# LOAD DATA
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("LOADING DATA")
+print("=" * 80)
+
+for _f in [FASTA_FILE, CSV_FILE, PAIRS_FILE]:
+    if not os.path.exists(_f):
+        print(f"ERROR: File not found: {_f}")
+        sys.exit(1)
+
+# --- 1. FASTA sequences ---
+print("Reading FASTA sequences...")
+fasta_seqs = {}          # {envhog_id: raw_aa_sequence}
+_cur_id  = None
+_cur_seq = []
+with open(FASTA_FILE) as fh:
+    for line in fh:
+        line = line.rstrip()
+        if line.startswith(">"):
+            if _cur_id is not None:
+                fasta_seqs[_cur_id] = "".join(_cur_seq)
+            _cur_id  = line[1:].split()[0]
+            _cur_seq = []
+        else:
+            _cur_seq.append(line)
+if _cur_id is not None:
+    fasta_seqs[_cur_id] = "".join(_cur_seq)
+print(f"  Sequences in FASTA: {len(fasta_seqs):,}")
+
+# --- 2. Protein metadata CSV ---
+print("Reading protein metadata CSV...")
+meta_df = pd.read_csv(CSV_FILE)
+print(f"  Loaded {len(meta_df):,} rows | columns: {meta_df.columns.tolist()}")
+
+# Keep only proteins that have a sequence in the FASTA
+meta_df = meta_df[meta_df["ENVHOG"].isin(fasta_seqs)].reset_index(drop=True)
+print(f"  After FASTA intersection: {len(meta_df):,} proteins retained")
+
+# --- 3. Build lookup structures ---
+envhog_to_viseq    = dict(zip(meta_df["ENVHOG"], meta_df["VISEQ"]))
+viseq_to_proteins  = defaultdict(list)   # viseq  → [envhog_id, ...]
+for row in meta_df.itertuples(index=False):
+    viseq_to_proteins[row.VISEQ].append(row.ENVHOG)
+
+# --- 4. Load positive VISEQ pairs ---
+print("Reading positive VISEQ pairs CSV...")
+pairs_df = pd.read_csv(PAIRS_FILE)
+print(f"  Loaded {len(pairs_df):,} positive VISEQ pairs")
+
+# Build bidirectional map: viseq → [list of positive viseqs]
+positive_viseq_map = defaultdict(list)
+for row in pairs_df.itertuples(index=False):
+    positive_viseq_map[row.viseq_A].append(row.viseq_B)
+    positive_viseq_map[row.viseq_B].append(row.viseq_A)
+print(f"  VISEQs with cross-VISEQ positives: {len(positive_viseq_map):,}")
+
+# Convert map values to sets for O(1) lookup during false-negative filtering
+positive_viseq_set = {k: set(v) for k, v in positive_viseq_map.items()}
+
+# ============================================================================
+# BUILD PROTEIN POOL
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("BUILDING PROTEIN POOL")
+print("=" * 80)
+
+
+def prepare_t5_seq(seq: str) -> str:
+    """Remove gaps, replace rare AAs, space-separate for T5 tokeniser."""
+    seq = seq.replace(" ", "")
+    seq = re.sub(r"[UZOB]", "X", seq)
+    return " ".join(list(seq))
+
+
+def sample_positive_protein(anchor_envhog: str, anchor_viseq: str):
+    """
+    Return one positive protein for the given anchor, or None if none exists.
+
+    Positive candidates:
+      1. Other proteins sharing the same VISEQ.
+      2. A protein from a cross-VISEQ positive (drawn from positive_viseq_map).
+
+    When both options are available, one is chosen at random to expose the model
+    to both types of similarity signal across the training epoch.
+    """
+    same_viseq  = [p for p in viseq_to_proteins[anchor_viseq] if p != anchor_envhog]
+    cross_viseqs = positive_viseq_map.get(anchor_viseq, [])
+
+    has_same  = bool(same_viseq)
+    has_cross = bool(cross_viseqs)
+
+    if not has_same and not has_cross:
+        return None   # orphan: no confirmed positive exists
+
+    if has_same and has_cross:
+        strategy = random.choice(["same", "cross"])
+    elif has_same:
+        strategy = "same"
+    else:
+        strategy = "cross"
+
+    if strategy == "same":
+        return random.choice(same_viseq)
+
+    # Cross-VISEQ: pick a random positive VISEQ, then a random protein from it
+    pos_viseq = random.choice(cross_viseqs)
+    pos_prots = viseq_to_proteins.get(pos_viseq, [])
+    if pos_prots:
+        return random.choice(pos_prots)
+    # The cross-VISEQ has no protein in our filtered dataset → fall back to same
+    if same_viseq:
+        return random.choice(same_viseq)
+    return None
+
+
+# Flat list of ALL proteins (orphans and non-orphans alike).
+# Positives are sampled on-the-fly in the collator; no pre-building of pairs needed.
+all_proteins = []  # list of (seq_t5, viseq, envhog_id)
+
+n_no_fasta = 0
+for row in meta_df.itertuples(index=False):
+    envhog  = row.ENVHOG
+    viseq   = row.VISEQ
+    raw_seq = fasta_seqs.get(envhog)
+    if raw_seq is None:
+        n_no_fasta += 1
+        continue
+    all_proteins.append((prepare_t5_seq(raw_seq), viseq, envhog))
+
+print(f"  Total proteins in pool:        {len(all_proteins):,}")
+print(f"  Skipped (no FASTA sequence):   {n_no_fasta:,}")
+print()
+print("  Orphan proteins automatically get adj_matrix rows of all-False:")
+print("  they contribute to MLM only, excluded from the Con mean.")
+
+if len(all_proteins) == 0:
+    print("ERROR: No proteins found. Check data paths.")
+    sys.exit(1)
+
+# ============================================================================
+# SPLIT DATA
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("SPLITTING DATA")
+print("=" * 80)
+
+random.shuffle(all_proteins)
+n_val_proteins = max(1, int(len(all_proteins) * 0.1))
+val_proteins   = all_proteins[:n_val_proteins]
+train_proteins = all_proteins[n_val_proteins:]
+
+if MAX_TRAIN_SAMPLES and len(train_proteins) > MAX_TRAIN_SAMPLES:
+    train_proteins = train_proteins[:MAX_TRAIN_SAMPLES]
+if MAX_EVAL_SAMPLES and len(val_proteins) > MAX_EVAL_SAMPLES:
+    val_proteins = val_proteins[:MAX_EVAL_SAMPLES]
+
+print(f"Training proteins:   {len(train_proteins):,}")
+print(f"Validation proteins: {len(val_proteins):,}")
+
+# ============================================================================
+# DATASET
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("BUILDING DATASET")
+print("=" * 80)
+
+
+class ProteinGraphDataset(torch.utils.data.Dataset):
+    """
+    A flat pool of all proteins. Each item is a single protein.
+    The PairGraphCollator samples positives on-the-fly and builds the
+    per-batch adjacency matrix for Contrastive loss.
+    """
+
+    def __init__(self, proteins):
+        # proteins: list of (seq_t5, viseq, envhog_id)
+        self.items = [
+            {"seq": s, "viseq": v, "envhog_id": e}
+            for s, v, e in proteins
+        ]
+
+    def __len__(self):
+        return len(self.items)
+
+    def __getitem__(self, idx):
+        return self.items[idx]
+
+
+train_dataset = ProteinGraphDataset(train_proteins)
+val_dataset   = ProteinGraphDataset(val_proteins)
+
+print(f"Train dataset: {len(train_dataset):,} proteins")
+print(f"Val dataset:   {len(val_dataset):,} proteins")
+print(f"Each batch of {BATCH_SIZE} drawn proteins → {BATCH_SIZE}–{2*BATCH_SIZE} unique proteins after positive sampling")
+
+# ============================================================================
+# LOAD TOKENIZER AND MODEL
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("LOADING TOKENIZER AND MODEL")
+print("=" * 80)
+
+tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, do_lower_case=False, legacy=True)
+print(f"Tokenizer loaded: {MODEL_NAME}")
+
+use_bf16     = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+model_dtype  = torch.bfloat16 if use_bf16 else torch.float32
+
+model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype=model_dtype)
+model.config.use_cache = False
+if hasattr(model, "gradient_checkpointing_enable"):
+    try:
+        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    except TypeError:
+        model.gradient_checkpointing_enable()
+
+print(f"Full Encoder-Decoder Model loaded: {MODEL_NAME}")
+print(f"Total parameters: {model.num_parameters():,}")
+
+# ============================================================================
+# CONFIGURE AND APPLY LORA
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("CONFIGURING LORA")
+print("=" * 80)
+
+lora_config = LoraConfig(
+    r=LORA_R,
+    lora_alpha=LORA_ALPHA,
+    target_modules=LORA_TARGET_MODULES,
+    lora_dropout=LORA_DROPOUT,
+    bias="none",
+    task_type=LORA_TASK_TYPE,
+)
+
+model = get_peft_model(model, lora_config)
+
+
+def _get_input_embedding_layer(model_obj):
+    getter = getattr(model_obj, "get_input_embeddings", None)
+    if callable(getter):
+        return getter()
+    base_model = getattr(model_obj, "base_model", None)
+    if base_model is not None:
+        base_getter = getattr(base_model, "get_input_embeddings", None)
+        if callable(base_getter):
+            return base_getter()
+    raise AttributeError("Could not resolve input embedding layer for model")
+
+
+def _make_inputs_require_grad(module, inputs, output):
+    output.requires_grad_(True)
+
+
+embedding_layer = cast(nn.Embedding, _get_input_embedding_layer(model))
+embedding_layer.register_forward_hook(_make_inputs_require_grad)
+
+model.print_trainable_parameters()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model  = model.to(device)
+print(f"Model moved to device: {device}")
+
+# ============================================================================
+# DATA COLLATOR
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("PREPARING DATA COLLATOR")
+print("=" * 80)
+
+# Build amino acid token IDs for random-replacement masking
+_aa_set = set()
+for _aa in "ACDEFGHIKLMNPQRSTVWY":
+    for _tid in tokenizer.encode(_aa, add_special_tokens=False):
+        if _tid not in (tokenizer.unk_token_id, tokenizer.eos_token_id, tokenizer.pad_token_id):
+            _aa_set.add(_tid)
+AA_TOKENS = list(_aa_set)
+if not AA_TOKENS:
+    AA_TOKENS = list(range(3, tokenizer.vocab_size))
+
+
+class PairGraphCollator:
+    """
+    Collates BATCH_SIZE drawn protein items into a graph-structured batch.
+
+    On-the-fly positive sampling:
+      For each drawn protein that has at least one known positive, sample 1
+      positive protein and add it to the pool (deduplicated by envhog_id).
+
+    Adjacency matrix A (N×N):
+      A[i,j] = True  iff  proteins i and j are known positives:
+                          same VISEQ  OR  cross-VISEQ positive in the pair graph.
+      Diagonal is always False (self-loops excluded).
+      Orphans (no positives anywhere) have all-False rows → contribute to MLM
+      only, automatically excluded from the Con mean by pos_count == 0.
+
+    Output keys:
+      all_input_ids      : (N, L) — BART-masked token ids
+      all_attention_mask : (N, L)
+      all_labels         : (N, L) — original tokens; padding → -100
+      adj_matrix         : Python list[list[bool]] (N×N), passed through as-is
+    """
+
+    def __init__(self, tokenizer, mlm_probability=0.15, pad_to_multiple_of=8):
+        self.tokenizer         = tokenizer
+        self.mlm_probability   = mlm_probability
+        self.pad_to_multiple_of = pad_to_multiple_of
+
+        mask_id = tokenizer.mask_token_id
+        if mask_id is None:
+            mask_id = tokenizer.convert_tokens_to_ids("<extra_id_0>")
+        self.mask_token_id = mask_id
+        self.pad_token_id  = tokenizer.pad_token_id
+        self.eos_token_id  = tokenizer.eos_token_id
+
+    def _apply_bart_mask(self, input_tensor, attention_tensor):
+        """Apply BART-style masking: 90% → <mask>, 10% → random amino acid."""
+        corrupted   = input_tensor.clone()
+        special_ids = {self.pad_token_id, self.eos_token_id}
+
+        prob_matrix = torch.full(input_tensor.shape, self.mlm_probability)
+        for sid in special_ids:
+            prob_matrix[input_tensor == sid] = 0.0
+        prob_matrix[attention_tensor == 0] = 0.0
+
+        mask_positions    = torch.bernoulli(prob_matrix).bool()
+        replace_with_mask = torch.bernoulli(
+            torch.full(mask_positions.shape, 0.9)
+        ).bool() & mask_positions
+        corrupted[replace_with_mask] = self.mask_token_id
+
+        replace_with_random = mask_positions & ~replace_with_mask
+        n_random = int(replace_with_random.sum().item())
+        if n_random > 0:
+            corrupted[replace_with_random] = torch.tensor(
+                random.choices(AA_TOKENS, k=n_random), dtype=torch.long
+            )
+        return corrupted
+
+    def __call__(self, features):
+        # features: list of BATCH_SIZE dicts {seq, viseq, envhog_id}
+
+        # Deduplicate drawn proteins by envhog_id (rare but possible)
+        pool_by_id = {}
+        for f in features:
+            eid = f["envhog_id"]
+            if eid not in pool_by_id:
+                pool_by_id[eid] = f
+
+        # For each drawn protein, sample 1 positive and add if not already in pool
+        for f in list(pool_by_id.values()):
+            pos_eid = sample_positive_protein(f["envhog_id"], f["viseq"])
+            if pos_eid is not None and pos_eid not in pool_by_id:
+                pos_viseq = envhog_to_viseq.get(pos_eid)
+                if pos_viseq is not None:
+                    pool_by_id[pos_eid] = {
+                        "seq":       prepare_t5_seq(fasta_seqs[pos_eid]),
+                        "viseq":     pos_viseq,
+                        "envhog_id": pos_eid,
+                    }
+
+        pool       = list(pool_by_id.values())  # N = 8..16 unique proteins
+        N          = len(pool)
+        all_seqs   = [p["seq"]   for p in pool]
+        all_viseqs = [p["viseq"] for p in pool]
+
+        # Build adjacency matrix (N×N) from the VISEQ pair graph
+        adj = []
+        for i in range(N):
+            vi      = all_viseqs[i]
+            pos_set = positive_viseq_set.get(vi, set()) | {vi}
+            adj.append([
+                (j != i and all_viseqs[j] in pos_set)
+                for j in range(N)
+            ])
+
+        # Tokenise all N sequences
+        encoding  = self.tokenizer(
+            all_seqs, truncation=True, max_length=MAX_LENGTH, add_special_tokens=True
+        )
+        all_ids   = encoding["input_ids"]
+        all_masks = encoding["attention_mask"]
+
+        # Pad to longest (aligned to pad_to_multiple_of)
+        max_len = max(len(ids) for ids in all_ids)
+        if self.pad_to_multiple_of:
+            max_len = (
+                (max_len + self.pad_to_multiple_of - 1)
+                // self.pad_to_multiple_of
+                * self.pad_to_multiple_of
+            )
+
+        pad_id       = self.pad_token_id
+        padded_ids   = []
+        padded_masks = []
+        for ids, mask in zip(all_ids, all_masks):
+            pad_len = max_len - len(ids)
+            padded_ids.append(ids + [pad_id] * pad_len)
+            padded_masks.append(mask + [0]   * pad_len)
+
+        input_tensor = torch.tensor(padded_ids,   dtype=torch.long)
+        attn_tensor  = torch.tensor(padded_masks, dtype=torch.long)
+
+        # Labels for MLM: original tokens; padding positions → -100
+        labels = input_tensor.clone()
+        labels[attn_tensor == 0] = -100
+
+        corrupted = self._apply_bart_mask(input_tensor, attn_tensor)
+
+        return {
+            "adj_matrix":         adj,          # Python list[list[bool]], passed through as-is
+            "all_input_ids":      corrupted,
+            "all_attention_mask": attn_tensor,
+            "all_labels":         labels,
+        }
+
+
+data_collator = PairGraphCollator(
+    tokenizer=tokenizer,
+    mlm_probability=NOISE_DENSITY,
+    pad_to_multiple_of=8,
+)
+print("Data collator: PairGraphCollator — on-the-fly positive sampling + Con adjacency matrix")
+print(f"  Each batch: {BATCH_SIZE} drawn proteins → {BATCH_SIZE}–{2*BATCH_SIZE} unique proteins after positive sampling")
+print("  Adjacency matrix built from VISEQ pair graph; orphan rows are all-False")
+
+# ============================================================================
+# CONTRASTIVE LOSS HELPERS
+# ============================================================================
+
+
+def mean_pool(hidden_states, attention_mask):
+    """
+    Mean-pool encoder last hidden states over non-padding token positions.
+    hidden_states  : (N, seq_len, hidden_dim)
+    attention_mask : (N, seq_len)   — 1 for real tokens, 0 for padding
+    returns        : (N, hidden_dim)
+    """
+    mask = attention_mask.unsqueeze(-1).float()
+    return (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+
+
+def con_loss(z, adj, temperature):
+    """
+    Contrastive Loss (Con) using a prebuilt adjacency matrix.
+
+    z   : (N, D) — unit-normalized encoder embeddings (F.normalize applied before calling)
+    adj : (N, N) bool — adj[i,j] = True means i and j are known positives
+    temperature : scalar
+
+    For each protein i that has at least one positive in the batch:
+      loss_i = -1/|P(i)| * sum_{j in P(i)} [sim(i,j)/T - log(sum_{k≠i} exp(sim(i,k)/T))]
+
+    Returns the mean over proteins with at least one positive.
+    Returns 0 (no grad) if no protein has a positive in the batch.
+    """
+    N   = z.size(0)
+    dev = z.device
+
+    # Full pairwise similarity matrix, temperature-scaled
+    sim = torch.matmul(z, z.T) / temperature               # (N, N)
+
+    # Mask diagonal so it does not contribute to the denominator
+    self_mask  = torch.eye(N, dtype=torch.bool, device=dev)
+    sim_masked = sim.masked_fill(self_mask, float("-inf"))
+
+    # log-sum-exp over all k≠i  →  log denominator for each anchor i
+    log_denom = torch.logsumexp(sim_masked, dim=1)         # (N,)
+
+    # log p(j | i) = sim[i,j]/T - log_denom[i]   for each j
+    log_prob  = sim - log_denom.unsqueeze(1)               # (N, N)
+
+    # Number of positives per protein
+    n_positives   = adj.float().sum(dim=1)                 # (N,)
+    has_positive  = n_positives > 0                        # (N,) bool
+
+    if not has_positive.any():
+        return torch.tensor(0.0, device=dev, requires_grad=True)
+
+    # Per-anchor loss: -1/|P(i)| * sum_{j: adj[i,j]} log_prob[i,j]
+    pos_log_sum   = (adj.float() * log_prob).sum(dim=1)    # (N,)
+    per_anchor    = -pos_log_sum / n_positives.clamp(min=1) # (N,)
+
+    return per_anchor[has_positive].mean()
+
+
+
+# ============================================================================
+# CURRICULUM LOSS (applied to MLM component — identical to Default_v2)
+# ============================================================================
+
+curriculum_state = {"global_step": 0, "max_steps": 1, "phase": "train"}
+
+# Shared state for component loss logging.
+# Written by ContraMLMTrainer.compute_loss at every training step;
+# read by ComponentLossLogCallback.on_log to inject into the Trainer log dict.
+_component_losses: dict = {"mlm_loss": None, "con_loss": None}
+
+
+def curriculum_keep_fraction(progress):
+    current_stage = int(np.floor(progress * NUM_STAGES))
+    if current_stage >= NUM_STAGES:
+        return KEEP_FRACTION_END
+    stage_size = (KEEP_FRACTION_END - KEEP_FRACTION_START) / (NUM_STAGES - 1)
+    return KEEP_FRACTION_START + (current_stage * stage_size)
+
+
+def curriculum_loss_from_outputs(outputs, labels):
+    logits = outputs.logits
+    token_losses = F.cross_entropy(
+        logits.view(-1, logits.size(-1)),
+        labels.view(-1),
+        ignore_index=-100,
+        reduction="none",
+    )
+    valid_tokens = labels.view(-1) != -100
+
+    if USE_LOSS_CLIPPING_CURRICULUM and torch.is_grad_enabled():
+        max_steps    = max(1, int(curriculum_state["max_steps"]))
+        progress     = min(1.0, float(curriculum_state["global_step"]) / float(max_steps))
+        keep_fraction = curriculum_keep_fraction(progress)
+
+        valid_indices = torch.nonzero(valid_tokens, as_tuple=False).squeeze(-1)
+        valid_losses  = token_losses[valid_tokens]
+        k             = max(1, int(valid_losses.numel() * keep_fraction))
+        selected_pos  = torch.topk(valid_losses, k=k, largest=LARGEST).indices
+        keep_indices  = valid_indices[selected_pos]
+        keep_tokens   = torch.zeros_like(valid_tokens, dtype=torch.bool)
+        keep_tokens[keep_indices] = True
+    else:
+        keep_tokens = valid_tokens
+
+    if not keep_tokens.any():
+        raise FloatingPointError("No valid tokens available for loss computation.")
+
+    return token_losses[keep_tokens].mean()
+
+
+# ============================================================================
+# CONFIGURE TRAINING ARGUMENTS
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("CONFIGURING TRAINING ARGUMENTS")
+print("=" * 80)
+
+dataloader_workers = min(8, os.cpu_count() or 1)
+
+training_kwargs = {
+    "output_dir":                    OUTPUT_DIR,
+    "save_strategy":                 "steps",
+    "eval_strategy":                 "steps",
+    "save_steps":                    1000,
+    "eval_steps":                    1000, # 5000
+    "gradient_accumulation_steps":   GRADIENT_ACCUMULATION_STEPS,
+    "per_device_train_batch_size":   BATCH_SIZE,
+    "per_device_eval_batch_size":    16,
+    "num_train_epochs":              NUM_EPOCHS,
+    "dataloader_num_workers":        dataloader_workers,
+    "dataloader_pin_memory":         True,
+    "logging_dir":                   f"{OUTPUT_DIR}/logs",
+    "logging_steps":                 100,
+    "save_total_limit":              3,
+    "fp16":                          False,
+    "bf16":                          use_bf16,
+    # REQUIRED: our batch dict uses custom keys (all_input_ids, etc.)
+    # that are not in the model's forward signature.
+    "remove_unused_columns":         False,
+    "load_best_model_at_end":        False,
+    "report_to":                     "none",
+    "push_to_hub":                   False,
+}
+
+training_signature = inspect.signature(TrainingArguments.__init__).parameters
+if "eval_strategy" not in training_signature:
+    training_kwargs.pop("eval_strategy", None)
+    training_kwargs["evaluation_strategy"] = "steps"
+
+if "bf16_full_eval" in training_signature:
+    training_kwargs["bf16_full_eval"] = use_bf16
+
+training_args = TrainingArguments(**training_kwargs)
+print("Training arguments configured")
+print(f"  remove_unused_columns = False  (required for custom batch keys)")
+print(
+    f"  Effective batch: {BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS} drawn proteins "
+    f"({BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS}–{2 * BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS} unique proteins per gradient step)"
+)
+
+# ============================================================================
+# CALLBACKS
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("INITIALIZING CALLBACKS")
+print("=" * 80)
+
+
+class LiveLossPlotCallback(TrainerCallback):
+    def __init__(self, output_dir):
+        self.output_dir = output_dir
+
+    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+        if metrics and "eval_loss" in metrics:
+            print(f"\n>>> [Step {state.global_step}] Evaluation Loss: {metrics['eval_loss']:.4f} <<<\n")
+        self._update_plot(state)
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        self._update_plot(state)
+
+    def _update_plot(self, state):
+        history     = state.log_history
+        train_loss  = [x["loss"]         for x in history if "loss"         in x]
+        train_steps = [x["step"]         for x in history if "loss"         in x]
+        eval_loss   = [x["eval_loss"]    for x in history if "eval_loss"    in x]
+        eval_steps  = [x["step"]         for x in history if "eval_loss"    in x]
+        mlm_loss    = [x["train_mlm_loss"]    for x in history if "train_mlm_loss"    in x]
+        con_loss = [x["train_con_loss"] for x in history if "train_con_loss" in x]
+        comp_steps  = [x["step"]         for x in history if "train_mlm_loss" in x]
+
+        if not train_loss:
+            return
+
+        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10), sharex=True)
+
+        # Top panel: total loss + eval loss
+        ax1.plot(train_steps, train_loss, label="Train Loss (total)",
+                 color="blue", alpha=0.6)
+        if eval_loss:
+            ax1.plot(eval_steps, eval_loss, label="Eval Loss (MLM only)",
+                     color="red", marker="o", linewidth=2)
+        ax1.set_ylabel("Loss")
+        ax1.grid(True, linestyle="--", alpha=0.6)
+        ax1.legend(loc="upper right")
+        ax1.set_title("Live Training Loss — ContraMLM v1")
+
+        # Bottom panel: MLM vs Con components
+        if mlm_loss:
+            ax2.plot(comp_steps, mlm_loss,    label="MLM loss",    color="green",  alpha=0.7)
+            ax2.plot(comp_steps, con_loss, label="Con loss", color="orange", alpha=0.7)
+            ax2.legend(loc="upper right")
+        ax2.set_xlabel("Training Steps")
+        ax2.set_ylabel("Component Loss")
+        ax2.grid(True, linestyle="--", alpha=0.6)
+
+        plt.tight_layout()
+        plt.savefig(os.path.join(self.output_dir, "live_loss_curve.png"), dpi=300)
+        plt.close()
+
+
+class CurriculumStateCallback(TrainerCallback):
+    def on_train_begin(self, args, state, control, **kwargs):
+        curriculum_state["global_step"] = state.global_step
+        curriculum_state["max_steps"]   = state.max_steps if state.max_steps and state.max_steps > 0 else 1
+        curriculum_state["phase"]       = "train"
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        curriculum_state["global_step"] = state.global_step
+        curriculum_state["phase"]       = "train"
+
+    def on_evaluate(self, args, state, control, **kwargs):
+        curriculum_state["phase"] = "eval"
+
+
+# ============================================================================
+# TRAINER
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("INITIALIZING TRAINER")
+print("=" * 80)
+
+
+class AdaFactorTrainer(Trainer):
+    def create_optimizer_and_scheduler(self, num_training_steps: int):
+        self.optimizer = Adafactor(
+            [p for p in self.model.parameters() if p.requires_grad],
+            scale_parameter=True,
+            relative_step=True,
+            warmup_init=True,
+            lr=None,
+        )
+        self.lr_scheduler = AdafactorSchedule(self.optimizer)
+
+
+class ContraMLMTrainer(AdaFactorTrainer):
+    """
+    Trainer combining BART-style MLM with Contrastive (Con) loss.
+
+    Batch layout (set by PairGraphCollator):
+        all_input_ids / all_attention_mask / all_labels : (N, L)
+        adj_matrix : Python list[list[bool]] (N×N)
+        N ranges from BATCH_SIZE (all orphans) to 2*BATCH_SIZE (all non-orphans).
+
+    Loss:
+        total = (1 - CONTRASTIVE_LAMBDA) * mlm_loss + CONTRASTIVE_LAMBDA * con_loss
+        Con is skipped (returns 0) if no protein in the batch has a positive.
+    """
+
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        # adj_matrix is not a model argument; remove it before the standard eval forward pass
+        inputs.pop("adj_matrix", None)
+        # Remap custom collator keys to the standard model argument names
+        if "all_input_ids" in inputs:
+            inputs["input_ids"] = inputs.pop("all_input_ids")
+        if "all_attention_mask" in inputs:
+            inputs["attention_mask"] = inputs.pop("all_attention_mask")
+        if "all_labels" in inputs:
+            inputs["labels"] = inputs.pop("all_labels")
+        return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # During training: custom keys (all_input_ids, etc.) + adj_matrix present.
+        # During eval:     prediction_step remaps keys to standard names, adj_matrix is gone.
+        adj_matrix         = inputs.get("adj_matrix", None)   # None during eval
+        all_input_ids      = inputs.get("all_input_ids",      inputs.get("input_ids"))
+        all_attention_mask = inputs.get("all_attention_mask", inputs.get("attention_mask"))
+        all_labels         = inputs.get("all_labels",         inputs.get("labels"))
+
+        # ---- Single encoder-decoder forward pass for all N proteins ----
+        outputs = model(
+            input_ids=all_input_ids,
+            attention_mask=all_attention_mask,
+            labels=all_labels,
+        )
+
+        # ---- MLM loss (curriculum-aware, same as Default_v2) ----
+        mlm_loss = curriculum_loss_from_outputs(outputs, all_labels)
+
+        # ---- Skip Con during eval (no adj_matrix) or all-orphan batch ----
+        if adj_matrix is None:
+            return (mlm_loss, outputs) if return_outputs else mlm_loss
+
+        # ---- Convert adjacency list to bool tensor ----
+        adj_tensor = torch.tensor(adj_matrix, dtype=torch.bool, device=all_input_ids.device)
+
+        if not adj_tensor.any():
+            _component_losses["mlm_loss"]    = mlm_loss.detach().float().item()
+            _component_losses["con_loss"] = 0.0
+            return (mlm_loss, outputs) if return_outputs else mlm_loss
+
+        # ---- Encoder embeddings → mean-pooled, unit-normalized representations ----
+        enc_hidden = outputs.encoder_last_hidden_state
+        z = mean_pool(enc_hidden, all_attention_mask)
+
+        # ---- Contrastive loss ----
+        contrastive_loss = con_loss(
+            F.normalize(z, dim=-1), adj_tensor, CONTRASTIVE_TEMPERATURE
+        )
+
+        _component_losses["mlm_loss"]    = mlm_loss.detach().float().item()
+        _component_losses["con_loss"] = contrastive_loss.detach().float().item()
+
+        total_loss = (1.0 - CONTRASTIVE_LAMBDA) * mlm_loss + CONTRASTIVE_LAMBDA * contrastive_loss
+        return (total_loss, outputs) if return_outputs else total_loss
+
+    def log(self, logs):
+        # Inject component losses into the log dict BEFORE the base class freezes
+        # it into state.log_history — this is the only way they appear in the history
+        # that _update_plot reads.
+        if _component_losses["mlm_loss"] is not None and "loss" in logs:
+            logs["train_mlm_loss"]    = round(_component_losses["mlm_loss"],    6)
+            logs["train_con_loss"] = round(_component_losses["con_loss"], 6)
+        super().log(logs)
+
+
+callbacks = [LiveLossPlotCallback(OUTPUT_DIR)]
+if USE_LOSS_CLIPPING_CURRICULUM:
+    callbacks.append(CurriculumStateCallback())
+
+trainer = ContraMLMTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=val_dataset,
+    data_collator=data_collator,
+    callbacks=callbacks,
+)
+
+print("ContraMLMTrainer initialised")
+print(f"  MLM loss weight:         {1.0 - CONTRASTIVE_LAMBDA:.2f}  (= 1 - CONTRASTIVE_LAMBDA)")
+print(f"  Con loss weight:      {CONTRASTIVE_LAMBDA}")
+print(f"  Contrastive temperature: {CONTRASTIVE_TEMPERATURE}")
+print(f"  Orphans: adj row all-False → auto-excluded from Con mean, MLM only")
+
+# ============================================================================
+# TRAIN THE MODEL
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("STARTING TRAINING")
+print("=" * 80)
+
+
+def find_latest_checkpoint(output_dir):
+    if not os.path.isdir(output_dir):
+        return None
+    latest_path = None
+    latest_step = -1
+    for entry in os.listdir(output_dir):
+        if not entry.startswith("checkpoint-"):
+            continue
+        step_str = entry.split("checkpoint-")[-1]
+        if not step_str.isdigit():
+            continue
+        full_path = os.path.join(output_dir, entry)
+        if not os.path.isdir(full_path):
+            continue
+        step = int(step_str)
+        if step > latest_step:
+            latest_step = step
+            latest_path = full_path
+    return latest_path
+
+
+def quarantine_rng_state_files(checkpoint_dir):
+    moved_files = []
+    for entry in os.listdir(checkpoint_dir):
+        if not (entry.startswith("rng_state") and entry.endswith(".pth")):
+            continue
+        src = os.path.join(checkpoint_dir, entry)
+        if not os.path.isfile(src):
+            continue
+        dst = src + ".bak"
+        os.replace(src, dst)
+        moved_files.append((src, dst))
+    return moved_files
+
+
+try:
+    resume_checkpoint = find_latest_checkpoint(OUTPUT_DIR)
+    if resume_checkpoint is not None:
+        print(f"Resuming training from checkpoint: {resume_checkpoint}")
+        moved_rng_files = quarantine_rng_state_files(resume_checkpoint)
+        if moved_rng_files:
+            print(
+                f"Skipped rng_state*.pth files for PyTorch 2.6 compatibility: "
+                f"{len(moved_rng_files)} file(s)."
+            )
+        train_result = trainer.train(resume_from_checkpoint=resume_checkpoint)
+    else:
+        print("No checkpoint found. Starting training from scratch.")
+        train_result = trainer.train()
+
+    print("\n" + "=" * 80)
+    print("TRAINING COMPLETED!")
+    print("=" * 80)
+    print(f"Train loss: {train_result.training_loss:.4f}")
+    print(f"Training time: {train_result.metrics['train_runtime']:.2f} seconds")
+
+except Exception as e:
+    print(f"\nERROR during training: {e}")
+    import traceback
+    traceback.print_exc()
+    sys.exit(1)
+
+# ============================================================================
+# EVALUATE THE MODEL
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("EVALUATING MODEL")
+print("=" * 80)
+
+try:
+    eval_results = trainer.evaluate()
+    print("Evaluation Results:")
+    for key, value in eval_results.items():
+        print(f"  {key}: {value:.4f}")
+except Exception as e:
+    print(f"ERROR during evaluation: {e}")
+
+# ============================================================================
+# SAVE THE MODEL
+# ============================================================================
+
+print("\n" + "=" * 80)
+print("SAVING MODEL")
+print("=" * 80)
+
+lora_output_dir = f"{OUTPUT_DIR}/lora_adapters"
+model.save_pretrained(lora_output_dir)
+tokenizer.save_pretrained(lora_output_dir)
+
+print(f"LoRA adapters saved to: {lora_output_dir}")
+
+print("\n" + "=" * 80)
+print("FINE-TUNING COMPLETE!")
+print("=" * 80)