Upload run_conformalesm_unified.py

run_conformalesm_unified.py

@@ -0,0 +1,672 @@
+"""
+ConformalESM-Extended: Unified Uncertainty Quantification for Protein Language Models
+=====================================================================================
+
+Runs overnight on Hugging Face CPU infrastructure. Pushes results to Hub.
+
+Coverage:
+  - Task 1: Secondary Structure (Q3: H/E/C) — 8M + 650M
+  - Task 2: Disorder Prediction (binary) — 35M + 650M
+  - Baselines: Entropy, Max-Margin
+  - Methods: Temperature Scaling, Split Conformal, Class-Conditional, Mondrian,
+            Size-Stratified, Experiment Prioritization, Cross-Model Transfer
+  - All post-hoc, CPU-only, no retraining
+
+Auto-pushes results to knoxel/conformalesm-paper-starter
+"""
+
+import os
+import sys
+import json
+import time
+import numpy as np
+from collections import defaultdict
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+import torch
+
+# ======================== CONFIG ========================
+
+SEED = 42
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+MAX_LEN = 1022
+N_CAL = 500
+N_TEST = 500
+
+# Models
+SS_MODEL_8M = "AmelieSchreiber/esm2_t6_8M_UR50D-finetuned-secondary-structure"
+SS_MODEL_650M = "gaodrew/esm2_t33_650M_UR50D-finetuned-secondary-structure"
+DIS_MODEL_35M = "CQSB/esm2_35M-LoRA-ID-DisProt7"
+DIS_MODEL_650M = "CQSB/esm2_650M-LoRA-ID-DisProt7"
+
+# Datasets
+SS_DATASET = "lamm-mit/protein_secondary_structure_from_PDB"
+DIS_DATASET = "CQSB/SoftDis"
+DIS_CONFIG = "id05"
+DIS_THRESHOLD = 0.5
+
+# Label mappings
+SS_ID2LABEL = {0: "C", 1: "H", 2: "E"}
+SS_LABEL2ID = {"C": 0, "H": 1, "E": 2}
+DIS_LABEL2ID = {"Ordered": 0, "Disordered": 1}
+
+# ======================== UTILS ========================
+
+def log(msg):
+    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
+
+def dssp_to_q3(c):
+    if c in "HGI":
+        return "H"
+    elif c in "EB":
+        return "E"
+    else:
+        return "C"
+
+# ======================== DATA ========================
+
+def load_ss_data():
+    ds = load_dataset(SS_DATASET, split="train")
+    ds = ds.filter(lambda x: x["Sequence_length"] <= MAX_LEN - 2)
+    ds = ds.shuffle(seed=SEED)
+    cal = ds.select(range(min(N_CAL, len(ds))))
+    test = ds.select(range(min(N_CAL, len(ds)), min(N_CAL + N_TEST, len(ds))))
+    return cal, test
+
+
+def load_disorder_data():
+    ds = load_dataset(DIS_DATASET, DIS_CONFIG)
+    train = ds["train"].shuffle(seed=SEED)
+    cal = train.select(range(min(N_CAL, len(train))))
+    test = ds["test"].shuffle(seed=SEED)
+    test = test.select(range(min(N_TEST, len(test))))
+    return cal, test
+
+# ======================== MODEL LOADING ========================
+
+def load_model(model_id, is_650m=False):
+    """Load model, handling LoRA adapters if needed."""
+    log(f"Loading model: {model_id}")
+
+    if "LoRA" in model_id or "lora" in model_id.lower():
+        from peft import PeftModel
+        if "650M" in model_id or "t33" in model_id:
+            base_id = "facebook/esm2_t33_650M_UR50D"
+        elif "35M" in model_id or "t12" in model_id:
+            base_id = "facebook/esm2_t12_35M_UR50D"
+        else:
+            base_id = "facebook/esm2_t6_8M_UR50D"
+
+        log(f"  Base model: {base_id}")
+        base = AutoModelForTokenClassification.from_pretrained(base_id)
+        model = PeftModel.from_pretrained(base, model_id)
+    else:
+        model = AutoModelForTokenClassification.from_pretrained(model_id)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model.eval()
+    log(f"  Model loaded. Parameters: {sum(p.numel() for p in model.parameters()):,}")
+    return model, tokenizer
+
+# ======================== INFERENCE ========================
+
+def infer_ss(model, tokenizer, dataset, batch_size=2):
+    """Secondary structure inference. Returns list of {true, probs, preds}."""
+    results = []
+    with torch.no_grad():
+        for i in range(0, len(dataset), batch_size):
+            batch = dataset[i:i + batch_size]
+            for j in range(len(batch["Sequence_spaced"])):
+                seq = batch["Sequence_spaced"][j].split()
+                ss = batch["Secondary_structure"][j][:len(seq)]
+                true = np.array([SS_LABEL2ID[dssp_to_q3(c)] for c in ss])
+
+                spaced = " ".join(seq[:MAX_LEN - 2])
+                inputs = tokenizer(spaced, return_tensors="pt", truncation=True, max_length=MAX_LEN)
+                logits = model(**inputs).logits.squeeze(0)
+                probs = torch.softmax(logits, dim=-1).numpy()
+
+                # Align to residues (skip special tokens)
+                input_ids = inputs["input_ids"].squeeze(0).tolist()
+                aligned_probs = []
+                residue_idx = 0
+                cls_id = tokenizer.cls_token_id
+                eos_id = tokenizer.eos_token_id
+                pad_id = tokenizer.pad_token_id
+
+                for tid in input_ids:
+                    if tid in [cls_id, eos_id, pad_id]:
+                        continue
+                    if residue_idx < len(true):
+                        aligned_probs.append(probs[residue_idx + 1])
+                        residue_idx += 1
+
+                aligned_probs = np.array(aligned_probs)
+                min_len = min(len(true), len(aligned_probs))
+                results.append({
+                    "true": true[:min_len],
+                    "probs": aligned_probs[:min_len],
+                    "preds": np.argmax(aligned_probs[:min_len], axis=-1),
+                })
+    return results
+
+
+def infer_disorder(model, tokenizer, dataset, batch_size=2):
+    """Disorder inference. Returns list of {true, probs, preds}."""
+    results = []
+    with torch.no_grad():
+        for i in range(0, len(dataset), batch_size):
+            batch = dataset[i:i + batch_size]
+            for j in range(len(batch["sequence"])):
+                seq = batch["sequence"][j]
+                freqs = batch["soft_disorder_frequency"][j]
+                true = np.array([1 if f >= DIS_THRESHOLD else 0 for f in freqs[:len(seq)]])
+
+                spaced = " ".join(list(seq)[:MAX_LEN - 2])
+                inputs = tokenizer(
+                    spaced, return_tensors="pt", truncation=True, max_length=MAX_LEN,
+                    return_special_tokens_mask=True
+                )
+                special_mask = inputs.pop("special_tokens_mask").squeeze(0).bool().numpy()
+
+                logits = model(**inputs).logits.squeeze(0)
+                probs = torch.softmax(logits, dim=-1).numpy()
+
+                # Remove special tokens
+                aligned_probs = probs[~special_mask]
+                min_len = min(len(true), len(aligned_probs))
+                results.append({
+                    "true": true[:min_len],
+                    "probs": aligned_probs[:min_len],
+                    "preds": np.argmax(aligned_probs[:min_len], axis=-1),
+                })
+    return results
+
+# ======================== METRICS ========================
+
+def compute_accuracy(results):
+    correct = sum(np.sum(r["preds"] == r["true"]) for r in results)
+    total = sum(len(r["true"]) for r in results)
+    return correct / total if total else 0
+
+
+def compute_ece(results, n_bins=10):
+    all_conf, all_correct = [], []
+    for r in results:
+        conf = np.max(r["probs"], axis=-1)
+        correct = (r["preds"] == r["true"]).astype(float)
+        all_conf.extend(conf)
+        all_correct.extend(correct)
+    all_conf = np.array(all_conf)
+    all_correct = np.array(all_correct)
+
+    ece_val = 0.0
+    for i in range(n_bins):
+        lo, hi = i / n_bins, (i + 1) / n_bins
+        mask = (all_conf > lo) & (all_conf <= hi)
+        if mask.sum() == 0:
+            continue
+        ece_val += mask.sum() * abs(all_conf[mask].mean() - all_correct[mask].mean())
+    return ece_val / len(all_conf) if len(all_conf) else 0
+
+
+def compute_brier(results):
+    scores = []
+    for r in results:
+        n = len(r["true"])
+        if n == 0:
+            continue
+        n_cls = r["probs"].shape[1]
+        one_hot = np.zeros((n, n_cls))
+        one_hot[np.arange(n), r["true"]] = 1
+        scores.append(np.mean(np.sum((r["probs"] - one_hot) ** 2, axis=-1)))
+    return np.mean(scores) if scores else 0
+
+# ======================== TEMPERATURE SCALING ========================
+
+def find_temperature(cal_results, grid=None):
+    if grid is None:
+        grid = np.linspace(0.5, 5.0, 50)
+    all_logits, all_labels = [], []
+    for r in cal_results:
+        probs = np.clip(r["probs"], 1e-10, 1.0)
+        all_logits.append(np.log(probs))
+        all_labels.append(r["true"])
+    all_logits = np.concatenate(all_logits)
+    all_labels = np.concatenate(all_labels)
+
+    best_t, best_nll = 1.0, float("inf")
+    for t in grid:
+        scaled = all_logits / t
+        max_log = np.max(scaled, axis=-1, keepdims=True)
+        log_probs = scaled - max_log - np.log(np.sum(np.exp(scaled - max_log), axis=-1, keepdims=True))
+        nll = -np.mean(log_probs[np.arange(len(all_labels)), all_labels])
+        if nll < best_nll:
+            best_nll = nll
+            best_t = t
+    return best_t
+
+
+def apply_temperature(results, temp):
+    scaled = []
+    for r in results:
+        probs = np.clip(r["probs"], 1e-10, 1.0)
+        logits = np.log(probs) / temp
+        max_log = np.max(logits, axis=-1, keepdims=True)
+        new_probs = np.exp(logits - max_log) / np.sum(np.exp(logits - max_log), axis=-1, keepdims=True)
+        scaled.append({"true": r["true"], "probs": new_probs, "preds": np.argmax(new_probs, axis=-1)})
+    return scaled
+
+# ======================== CONFORMAL ========================
+
+def conformal_qhat(cal_results, alpha=0.1):
+    scores = [1.0 - r["probs"][j, label]
+              for r in cal_results for j, label in enumerate(r["true"])]
+    scores = np.array(scores)
+    n = len(scores)
+    q = np.ceil((n + 1) * (1 - alpha)) / n
+    return np.quantile(scores, q, method="higher")
+
+
+def conformal_qhat_class_conditional(cal_results, alpha=0.1):
+    class_scores = defaultdict(list)
+    for r in cal_results:
+        for j, label in enumerate(r["true"]):
+            class_scores[label].append(1.0 - r["probs"][j, label])
+    thresholds = {}
+    for label, scores in class_scores.items():
+        scores = np.array(scores)
+        n = len(scores)
+        if n == 0:
+            thresholds[label] = 1.0
+            continue
+        q = np.ceil((n + 1) * (1 - alpha)) / n
+        thresholds[label] = np.quantile(scores, q, method="higher")
+    return thresholds
+
+
+def evaluate_conformal(results, q_hat, n_classes, per_class_thresholds=None):
+    coverage_count, total = 0, 0
+    set_sizes = []
+    class_cov = defaultdict(int)
+    class_tot = defaultdict(int)
+    class_set = defaultdict(list)
+    size_strat = defaultdict(lambda: {"correct": 0, "total": 0})
+
+    for r in results:
+        for j, label in enumerate(r["true"]):
+            total += 1
+            threshold = per_class_thresholds.get(label, q_hat) if per_class_thresholds else q_hat
+            pred_set = [y for y in range(n_classes) if (1.0 - r["probs"][j, y]) <= threshold]
+            set_size = len(pred_set)
+            set_sizes.append(set_size)
+            size_strat[set_size]["total"] += 1
+
+            if label in pred_set:
+                coverage_count += 1
+                class_cov[label] += 1
+                size_strat[set_size]["correct"] += 1
+            class_tot[label] += 1
+            class_set[label].append(set_size)
+
+    coverage = coverage_count / total if total else 0
+    avg_size = np.mean(set_sizes) if set_sizes else 0
+
+    per_class = {}
+    for k in sorted(class_tot.keys()):
+        per_class[k] = {
+            "coverage": class_cov[k] / class_tot[k] if class_tot[k] else 0,
+            "avg_set_size": np.mean(class_set[k]) if class_set[k] else 0,
+        }
+
+    size_strat_out = {}
+    for size in sorted(size_strat.keys()):
+        d = size_strat[size]
+        size_strat_out[size] = {
+            "coverage": d["correct"] / d["total"] if d["total"] else 0,
+            "n": d["total"],
+        }
+
+    return coverage, avg_size, per_class, size_strat_out
+
+
+def evaluate_mondrian(cal_results, test_results, alpha, n_classes):
+    class_cal = defaultdict(list)
+    for r in cal_results:
+        for j, label in enumerate(r["true"]):
+            class_cal[label].append(1.0 - r["probs"][j, label])
+
+    thresholds = {}
+    for label, scores in class_cal.items():
+        scores = np.array(scores)
+        n = len(scores)
+        if n == 0:
+            thresholds[label] = 1.0
+            continue
+        q = np.ceil((n + 1) * (1 - alpha)) / n
+        thresholds[label] = np.quantile(scores, q, method="higher")
+
+    class_cov = defaultdict(lambda: {"correct": 0, "total": 0})
+    class_set = defaultdict(list)
+
+    for r in test_results:
+        for j, label in enumerate(r["true"]):
+            threshold = thresholds.get(label, 1.0)
+            pred_set = [y for y in range(n_classes) if (1.0 - r["probs"][j, y]) <= threshold]
+            set_size = len(pred_set)
+            class_cov[label]["total"] += 1
+            class_set[label].append(set_size)
+            if label in pred_set:
+                class_cov[label]["correct"] += 1
+
+    mondrian = {}
+    for k in sorted(class_cov.keys()):
+        d = class_cov[k]
+        mondrian[k] = {
+            "coverage": d["correct"] / d["total"] if d["total"] else 0,
+            "avg_set_size": np.mean(class_set[k]) if class_set[k] else 0,
+            "n": d["total"],
+        }
+    return mondrian
+
+# ======================== BASELINES ========================
+
+def entropy_baseline(results, alpha, n_classes):
+    coverage_count, total = 0, 0
+    set_sizes = []
+    for r in results:
+        for j, label in enumerate(r["true"]):
+            total += 1
+            probs = r["probs"][j]
+            sorted_idx = np.argsort(-probs)
+            cumsum = np.cumsum(probs[sorted_idx])
+            n_include = np.searchsorted(cumsum, 1 - alpha) + 1
+            pred_set = sorted_idx[:n_include].tolist()
+            set_sizes.append(len(pred_set))
+            if label in pred_set:
+                coverage_count += 1
+    return coverage_count / total if total else 0, np.mean(set_sizes) if set_sizes else 0
+
+
+def maxmargin_baseline(results, alpha, n_classes):
+    all_margins = []
+    for r in results:
+        for j in range(len(r["true"])):
+            probs = r["probs"][j]
+            sp = np.sort(probs)[::-1]
+            all_margins.append(sp[0] - sp[1] if len(sp) > 1 else 1.0)
+    all_margins = np.array(all_margins)
+
+    n = len(all_margins)
+    q = np.ceil((n + 1) * (1 - alpha)) / n
+    margin_thresh = np.quantile(all_margins, q, method="higher")
+
+    coverage_count, total = 0, 0
+    set_sizes = []
+    for r in results:
+        for j, label in enumerate(r["true"]):
+            total += 1
+            probs = r["probs"][j]
+            sorted_idx = np.argsort(-probs)
+            sp = np.sort(probs)[::-1]
+            margin = sp[0] - sp[1] if len(sp) > 1 else 1.0
+            if margin >= margin_thresh:
+                pred_set = [sorted_idx[0]]
+            else:
+                pred_set = sorted_idx[:min(2, n_classes)].tolist()
+            set_sizes.append(len(pred_set))
+            if label in pred_set:
+                coverage_count += 1
+    return coverage_count / total if total else 0, np.mean(set_sizes) if set_sizes else 0
+
+# ======================== PRIORITIZATION ========================
+
+def experiment_prioritization(results, budgets):
+    all_unc, all_errors = [], []
+    for r in results:
+        max_probs = np.max(r["probs"], axis=-1)
+        uncertainties = 1 - max_probs
+        errors = (r["preds"] != r["true"]).astype(float)
+        all_unc.extend(uncertainties)
+        all_errors.extend(errors)
+
+    all_unc = np.array(all_unc)
+    all_errors = np.array(all_errors)
+    n_total = len(all_unc)
+
+    out = {}
+    for budget in budgets:
+        b = min(budget, n_total)
+        random_idx = np.random.choice(n_total, size=b, replace=False)
+        random_rate = all_errors[random_idx].mean()
+
+        sorted_idx = np.argsort(-all_unc)
+        top_idx = sorted_idx[:b]
+        unc_rate = all_errors[top_idx].mean()
+
+        catch = unc_rate / random_rate if random_rate > 0 else float('inf')
+        out[budget] = {
+            "random_error_rate": float(random_rate),
+            "uncertainty_error_rate": float(unc_rate),
+            "catch_rate": float(catch),
+        }
+    return out
+
+# ======================== CROSS-MODEL ========================
+
+def cross_model_transfer(cal_results_small, test_results_large, alpha, n_classes):
+    q = conformal_qhat(cal_results_small, alpha)
+    cov, size, _, size_strat = evaluate_conformal(test_results_large, q, n_classes)
+    return {
+        "q_hat": float(q),
+        "coverage": float(cov),
+        "avg_set_size": float(size),
+        "size_stratified": {str(k): v for k, v in size_strat.items()},
+    }
+
+# ======================== PIPELINE ========================
+
+def run_pipeline(model_id, dataset_loader, infer_fn, task_name, n_classes, label_map,
+                 budgets=[100, 500, 1000, 5000]):
+    """Run full conformal pipeline for one model+task."""
+    log(f"\n{'='*60}")
+    log(f"TASK: {task_name}")
+    log(f"MODEL: {model_id}")
+    log(f"{'='*60}")
+
+    model, tokenizer = load_model(model_id)
+    cal_ds, test_ds = dataset_loader()
+    log(f"  Calibration: {len(cal_ds)} seqs, Test: {len(test_ds)} seqs")
+
+    log("  Running inference (calibration)...")
+    cal_results = infer_fn(model, tokenizer, cal_ds)
+    log(f"  Calibration residues: {sum(len(r['true']) for r in cal_results):,}")
+
+    log("  Running inference (test)...")
+    test_results = infer_fn(model, tokenizer, test_ds)
+    log(f"  Test residues: {sum(len(r['true']) for r in test_results):,}")
+
+    # Free model memory
+    del model
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    # Baseline
+    base_acc = compute_accuracy(test_results)
+    base_ece = compute_ece(test_results)
+    base_brier = compute_brier(test_results)
+    log(f"  Baseline: Acc={base_acc:.4f}, ECE={base_ece:.4f}, Brier={base_brier:.4f}")
+
+    # Temperature scaling
+    best_t = find_temperature(cal_results)
+    scaled_cal = apply_temperature(cal_results, best_t)
+    scaled_test = apply_temperature(test_results, best_t)
+    ts_acc = compute_accuracy(scaled_test)
+    ts_ece = compute_ece(scaled_test)
+    ts_brier = compute_brier(scaled_test)
+    ece_red = (base_ece - ts_ece) / base_ece * 100 if base_ece else 0
+    log(f"  Temperature T={best_t:.2f}: Acc={ts_acc:.4f}, ECE={ts_ece:.4f} ({ece_red:+.0f}%), Brier={ts_brier:.4f}")
+
+    # Conformal (raw)
+    log("  Conformal prediction...")
+    conformal = {}
+    for alpha in [0.05, 0.10, 0.20]:
+        q = conformal_qhat(cal_results, alpha)
+        cov, size, pclass, sstrat = evaluate_conformal(test_results, q, n_classes)
+        log(f"    Raw alpha={alpha:.2f}: cov={cov:.4f}, set={size:.2f}")
+
+        q_s = conformal_qhat(scaled_cal, alpha)
+        cov_s, size_s, pclass_s, sstrat_s = evaluate_conformal(scaled_test, q_s, n_classes)
+        log(f"    T-scaled alpha={alpha:.2f}: cov={cov_s:.4f}, set={size_s:.2f}")
+
+        conformal[f"alpha_{alpha}"] = {
+            "raw": {"coverage": float(cov), "avg_set_size": float(size),
+                    "per_class": {label_map.get(k, str(k)): v for k, v in pclass.items()},
+                    "size_stratified": {str(kk): vv for kk, vv in sstrat.items()}},
+            "temperature_scaled": {"coverage": float(cov_s), "avg_set_size": float(size_s),
+                                   "per_class": {label_map.get(k, str(k)): v for k, v in pclass_s.items()},
+                                   "size_stratified": {str(kk): vv for kk, vv in sstrat_s.items()}},
+        }
+
+    # Class-conditional
+    log("  Class-conditional conformal...")
+    cc = {}
+    for alpha in [0.05, 0.10, 0.20]:
+        th = conformal_qhat_class_conditional(cal_results, alpha)
+        cov, size, pclass, _ = evaluate_conformal(test_results, 0, n_classes, th)
+        log(f"    alpha={alpha:.2f}: cov={cov:.4f}, set={size:.2f}")
+        cc[f"alpha_{alpha}"] = {
+            "coverage": float(cov), "avg_set_size": float(size),
+            "per_class": {label_map.get(k, str(k)): v for k, v in pclass.items()},
+        }
+
+    # Mondrian
+    log("  Mondrian conformal...")
+    mondrian = {}
+    for alpha in [0.05, 0.10, 0.20]:
+        mon = evaluate_mondrian(cal_results, test_results, alpha, n_classes)
+        log(f"    alpha={alpha:.2f}")
+        for k, v in mon.items():
+            log(f"      {label_map.get(k, str(k))}: cov={v['coverage']:.4f}, set={v['avg_set_size']:.2f}, n={v['n']}")
+        mondrian[f"alpha_{alpha}"] = {label_map.get(k, str(k)): v for k, v in mon.items()}
+
+    # Baselines
+    log("  Baselines...")
+    ent = {}
+    mm = {}
+    for alpha in [0.05, 0.10, 0.20]:
+        ec, es = entropy_baseline(test_results, alpha, n_classes)
+        mc, ms = maxmargin_baseline(test_results, alpha, n_classes)
+        log(f"    alpha={alpha:.2f}: Entropy cov={ec:.4f} set={es:.2f}, MaxMargin cov={mc:.4f} set={ms:.2f}")
+        ent[f"alpha_{alpha}"] = {"coverage": float(ec), "avg_set_size": float(es)}
+        mm[f"alpha_{alpha}"] = {"coverage": float(mc), "avg_set_size": float(ms)}
+
+    # Prioritization
+    log("  Experiment prioritization...")
+    prio = experiment_prioritization(test_results, budgets)
+    for b, d in prio.items():
+        log(f"    Budget={b}: random={d['random_error_rate']:.3f}, unc={d['uncertainty_error_rate']:.3f}, catch={d['catch_rate']:.2f}x")
+
+    return {
+        "task": task_name,
+        "model": model_id,
+        "baseline": {"accuracy": float(base_acc), "ece": float(base_ece), "brier": float(base_brier)},
+        "temperature_scaling": {"temperature": float(best_t), "accuracy": float(ts_acc),
+                                 "ece": float(ts_ece), "brier": float(ts_brier),
+                                 "ece_reduction_pct": float(ece_red)},
+        "conformal": conformal,
+        "class_conditional": cc,
+        "mondrian": mondrian,
+        "entropy_baseline": ent,
+        "maxmargin_baseline": mm,
+        "experiment_prioritization": prio,
+        "_cal_raw": cal_results,
+        "_cal_scaled": scaled_cal,
+        "_test_raw": test_results,
+        "_test_scaled": scaled_test,
+    }
+
+# ======================== MAIN ========================
+
+def main():
+    log("=" * 60)
+    log("ConformalESM-Extended: Unified Pipeline")
+    log("All post-hoc, CPU-only, no retraining")
+    log("=" * 60)
+
+    all_results = {}
+
+    # Task 1: Secondary Structure - 8M
+    ss8m = run_pipeline(SS_MODEL_8M, load_ss_data, infer_ss,
+                        "Secondary Structure (Q3) - ESM-2-8M", 3, SS_ID2LABEL)
+    all_results["ss_8m"] = {k: v for k, v in ss8m.items() if not k.startswith("_")}
+
+    # Task 1: Secondary Structure - 650M
+    ss650m = run_pipeline(SS_MODEL_650M, load_ss_data, infer_ss,
+                          "Secondary Structure (Q3) - ESM-2-650M", 3, SS_ID2LABEL)
+    all_results["ss_650m"] = {k: v for k, v in ss650m.items() if not k.startswith("_")}
+
+    # Task 2: Disorder - 35M
+    dis35m = run_pipeline(DIS_MODEL_35M, load_disorder_data, infer_disorder,
+                          "Disorder Prediction - ESM-2-35M", 2, {0: "Ordered", 1: "Disordered"})
+    all_results["disorder_35m"] = {k: v for k, v in dis35m.items() if not k.startswith("_")}
+
+    # Task 2: Disorder - 650M
+    dis650m = run_pipeline(DIS_MODEL_650M, load_disorder_data, infer_disorder,
+                           "Disorder Prediction - ESM-2-650M", 2, {0: "Ordered", 1: "Disordered"})
+    all_results["disorder_650m"] = {k: v for k, v in dis650m.items() if not k.startswith("_")}
+
+    # Cross-model transfer
+    log(f"\n{'='*60}")
+    log("Cross-Model Calibration Transfer")
+    log(f"{'='*60}")
+
+    log("  SS: 8M calibrate -> 650M test...")
+    t1 = cross_model_transfer(ss8m["_cal_raw"], ss650m["_test_raw"], 0.10, 3)
+    log(f"    Coverage: {t1['coverage']:.4f}, Avg set: {t1['avg_set_size']:.2f}")
+    all_results["transfer_ss_8m_to_650m"] = t1
+
+    log("  SS: 8M calibrate (T-scaled) -> 650M test...")
+    t2 = cross_model_transfer(ss8m["_cal_scaled"], ss650m["_test_scaled"], 0.10, 3)
+    log(f"    Coverage: {t2['coverage']:.4f}, Avg set: {t2['avg_set_size']:.2f}")
+    all_results["transfer_ss_8m_to_650m_temperature_scaled"] = t2
+
+    log("  Disorder: 35M calibrate -> 650M test...")
+    t3 = cross_model_transfer(dis35m["_cal_raw"], dis650m["_test_raw"], 0.10, 2)
+    log(f"    Coverage: {t3['coverage']:.4f}, Avg set: {t3['avg_set_size']:.2f}")
+    all_results["transfer_dis_35m_to_650m"] = t3
+
+    # Save
+    log(f"\n{'='*60}")
+    log("Saving Results")
+    log(f"{'='*60}")
+
+    with open("conformalesm_unified_results.json", "w") as f:
+        json.dump(all_results, f, indent=2)
+    log("  Saved: conformalesm_unified_results.json")
+
+    # Push to hub
+    log("  Pushing to knoxel/conformalesm-paper-starter...")
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj="conformalesm_unified_results.json",
+            path_in_repo="conformalesm_unified_results.json",
+            repo_id="knoxel/conformalesm-paper-starter",
+            repo_type="model",
+        )
+        log("  Successfully pushed to Hub!")
+    except Exception as e:
+        log(f"  Could not push to Hub: {e}")
+        log("  Results saved locally.")
+
+    log(f"\n{'='*60}")
+    log("ALL DONE")
+    log(f"{'='*60}")
+
+
+if __name__ == "__main__":
+    main()