Transformers
English
conformal-prediction
protein-language-models
uncertainty-quantification
esm-2
temperature-scaling
cpu
protein-structure
protein-engineering
Instructions to use knoxel/conformalesm-paper-starter with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use knoxel/conformalesm-paper-starter with Transformers:
# Load model directly from transformers import AutoModel model = AutoModel.from_pretrained("knoxel/conformalesm-paper-starter", device_map="auto") - Notebooks
- Google Colab
- Kaggle
Upload run_conformalesm_job1.py
Browse files- run_conformalesm_job1.py +563 -0
run_conformalesm_job1.py
ADDED
|
@@ -0,0 +1,563 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
ConformalESM Job 1: Core Experiments (CPU-only)
|
| 3 |
+
- Secondary Structure: ESM-2-8M
|
| 4 |
+
- Disorder: ESM-2-35M
|
| 5 |
+
+ All baselines, temperature scaling, conformal variants, experiment prioritization
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import os
|
| 9 |
+
import json
|
| 10 |
+
import time
|
| 11 |
+
import numpy as np
|
| 12 |
+
from collections import defaultdict
|
| 13 |
+
from datasets import load_dataset
|
| 14 |
+
from transformers import AutoTokenizer, AutoModelForTokenClassification
|
| 15 |
+
import torch
|
| 16 |
+
|
| 17 |
+
SEED = 42
|
| 18 |
+
np.random.seed(SEED)
|
| 19 |
+
torch.manual_seed(SEED)
|
| 20 |
+
|
| 21 |
+
MAX_LEN = 1022
|
| 22 |
+
N_CAL = 500
|
| 23 |
+
N_TEST = 500
|
| 24 |
+
|
| 25 |
+
# Models
|
| 26 |
+
SS_MODEL = "AmelieSchreiber/esm2_t6_8M_UR50D-finetuned-secondary-structure"
|
| 27 |
+
DIS_MODEL = "CQSB/esm2_35M-LoRA-ID-DisProt7"
|
| 28 |
+
|
| 29 |
+
# Datasets
|
| 30 |
+
SS_DATASET = "lamm-mit/protein_secondary_structure_from_PDB"
|
| 31 |
+
DIS_DATASET = "CQSB/SoftDis"
|
| 32 |
+
DIS_CONFIG = "id05"
|
| 33 |
+
DIS_THRESHOLD = 0.5
|
| 34 |
+
|
| 35 |
+
SS_ID2LABEL = {0: "C", 1: "H", 2: "E"}
|
| 36 |
+
SS_LABEL2ID = {"C": 0, "H": 1, "E": 2}
|
| 37 |
+
|
| 38 |
+
def log(msg):
|
| 39 |
+
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
| 40 |
+
|
| 41 |
+
def dssp_to_q3(c):
|
| 42 |
+
if c in "HGI": return "H"
|
| 43 |
+
elif c in "EB": return "E"
|
| 44 |
+
else: return "C"
|
| 45 |
+
|
| 46 |
+
# ===================== DATA =====================
|
| 47 |
+
|
| 48 |
+
def load_ss_data():
|
| 49 |
+
ds = load_dataset(SS_DATASET, split="train")
|
| 50 |
+
ds = ds.filter(lambda x: x["Sequence_length"] <= MAX_LEN - 2)
|
| 51 |
+
ds = ds.shuffle(seed=SEED)
|
| 52 |
+
cal = ds.select(range(min(N_CAL, len(ds))))
|
| 53 |
+
test = ds.select(range(min(N_CAL, len(ds)), min(N_CAL + N_TEST, len(ds))))
|
| 54 |
+
return cal, test
|
| 55 |
+
|
| 56 |
+
def load_disorder_data():
|
| 57 |
+
ds = load_dataset(DIS_DATASET, DIS_CONFIG)
|
| 58 |
+
train = ds["train"].shuffle(seed=SEED)
|
| 59 |
+
cal = train.select(range(min(N_CAL, len(train))))
|
| 60 |
+
test = ds["test"].shuffle(seed=SEED)
|
| 61 |
+
test = test.select(range(min(N_TEST, len(test))))
|
| 62 |
+
return cal, test
|
| 63 |
+
|
| 64 |
+
# ===================== MODEL =====================
|
| 65 |
+
|
| 66 |
+
def load_model(model_id):
|
| 67 |
+
log(f"Loading model: {model_id}")
|
| 68 |
+
if "LoRA" in model_id or "lora" in model_id.lower():
|
| 69 |
+
from peft import PeftModel
|
| 70 |
+
if "35M" in model_id or "t12" in model_id:
|
| 71 |
+
base_id = "facebook/esm2_t12_35M_UR50D"
|
| 72 |
+
elif "650M" in model_id or "t33" in model_id:
|
| 73 |
+
base_id = "facebook/esm2_t33_650M_UR50D"
|
| 74 |
+
else:
|
| 75 |
+
base_id = "facebook/esm2_t6_8M_UR50D"
|
| 76 |
+
base = AutoModelForTokenClassification.from_pretrained(base_id)
|
| 77 |
+
model = PeftModel.from_pretrained(base, model_id)
|
| 78 |
+
else:
|
| 79 |
+
model = AutoModelForTokenClassification.from_pretrained(model_id)
|
| 80 |
+
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
| 81 |
+
model.eval()
|
| 82 |
+
log(f" Model loaded. Params: {sum(p.numel() for p in model.parameters()):,}")
|
| 83 |
+
return model, tokenizer
|
| 84 |
+
|
| 85 |
+
# ===================== INFERENCE =====================
|
| 86 |
+
|
| 87 |
+
def infer_ss(model, tokenizer, dataset, batch_size=2):
|
| 88 |
+
results = []
|
| 89 |
+
with torch.no_grad():
|
| 90 |
+
for i in range(0, len(dataset), batch_size):
|
| 91 |
+
batch = dataset[i:i + batch_size]
|
| 92 |
+
for j in range(len(batch["Sequence_spaced"])):
|
| 93 |
+
seq = batch["Sequence_spaced"][j].split()
|
| 94 |
+
ss = batch["Secondary_structure"][j][:len(seq)]
|
| 95 |
+
true = np.array([SS_LABEL2ID[dssp_to_q3(c)] for c in ss])
|
| 96 |
+
spaced = " ".join(seq[:MAX_LEN - 2])
|
| 97 |
+
inputs = tokenizer(spaced, return_tensors="pt", truncation=True, max_length=MAX_LEN)
|
| 98 |
+
logits = model(**inputs).logits.squeeze(0)
|
| 99 |
+
probs = torch.softmax(logits, dim=-1).numpy()
|
| 100 |
+
input_ids = inputs["input_ids"].squeeze(0).tolist()
|
| 101 |
+
aligned_probs = []
|
| 102 |
+
residue_idx = 0
|
| 103 |
+
cls_id = tokenizer.cls_token_id
|
| 104 |
+
eos_id = tokenizer.eos_token_id
|
| 105 |
+
pad_id = tokenizer.pad_token_id
|
| 106 |
+
for tid in input_ids:
|
| 107 |
+
if tid in [cls_id, eos_id, pad_id]:
|
| 108 |
+
continue
|
| 109 |
+
if residue_idx < len(true):
|
| 110 |
+
aligned_probs.append(probs[residue_idx + 1])
|
| 111 |
+
residue_idx += 1
|
| 112 |
+
aligned_probs = np.array(aligned_probs)
|
| 113 |
+
min_len = min(len(true), len(aligned_probs))
|
| 114 |
+
results.append({
|
| 115 |
+
"true": true[:min_len],
|
| 116 |
+
"probs": aligned_probs[:min_len],
|
| 117 |
+
"preds": np.argmax(aligned_probs[:min_len], axis=-1),
|
| 118 |
+
})
|
| 119 |
+
return results
|
| 120 |
+
|
| 121 |
+
def infer_disorder(model, tokenizer, dataset, batch_size=2):
|
| 122 |
+
results = []
|
| 123 |
+
with torch.no_grad():
|
| 124 |
+
for i in range(0, len(dataset), batch_size):
|
| 125 |
+
batch = dataset[i:i + batch_size]
|
| 126 |
+
for j in range(len(batch["sequence"])):
|
| 127 |
+
seq = batch["sequence"][j]
|
| 128 |
+
freqs = batch["soft_disorder_frequency"][j]
|
| 129 |
+
true = np.array([1 if f >= DIS_THRESHOLD else 0 for f in freqs[:len(seq)]])
|
| 130 |
+
spaced = " ".join(list(seq)[:MAX_LEN - 2])
|
| 131 |
+
inputs = tokenizer(spaced, return_tensors="pt", truncation=True, max_length=MAX_LEN, return_special_tokens_mask=True)
|
| 132 |
+
special_mask = inputs.pop("special_tokens_mask").squeeze(0).bool().numpy()
|
| 133 |
+
logits = model(**inputs).logits.squeeze(0)
|
| 134 |
+
probs = torch.softmax(logits, dim=-1).numpy()
|
| 135 |
+
aligned_probs = probs[~special_mask]
|
| 136 |
+
min_len = min(len(true), len(aligned_probs))
|
| 137 |
+
results.append({
|
| 138 |
+
"true": true[:min_len],
|
| 139 |
+
"probs": aligned_probs[:min_len],
|
| 140 |
+
"preds": np.argmax(aligned_probs[:min_len], axis=-1),
|
| 141 |
+
})
|
| 142 |
+
return results
|
| 143 |
+
|
| 144 |
+
# ===================== METRICS =====================
|
| 145 |
+
|
| 146 |
+
def compute_accuracy(results):
|
| 147 |
+
correct = sum(np.sum(r["preds"] == r["true"]) for r in results)
|
| 148 |
+
total = sum(len(r["true"]) for r in results)
|
| 149 |
+
return correct / total if total else 0
|
| 150 |
+
|
| 151 |
+
def compute_ece(results, n_bins=10):
|
| 152 |
+
all_conf, all_correct = [], []
|
| 153 |
+
for r in results:
|
| 154 |
+
conf = np.max(r["probs"], axis=-1)
|
| 155 |
+
correct = (r["preds"] == r["true"]).astype(float)
|
| 156 |
+
all_conf.extend(conf)
|
| 157 |
+
all_correct.extend(correct)
|
| 158 |
+
all_conf = np.array(all_conf)
|
| 159 |
+
all_correct = np.array(all_correct)
|
| 160 |
+
ece_val = 0.0
|
| 161 |
+
for i in range(n_bins):
|
| 162 |
+
lo, hi = i / n_bins, (i + 1) / n_bins
|
| 163 |
+
mask = (all_conf > lo) & (all_conf <= hi)
|
| 164 |
+
if mask.sum() == 0: continue
|
| 165 |
+
ece_val += mask.sum() * abs(all_conf[mask].mean() - all_correct[mask].mean())
|
| 166 |
+
return ece_val / len(all_conf) if len(all_conf) else 0
|
| 167 |
+
|
| 168 |
+
def compute_brier(results):
|
| 169 |
+
scores = []
|
| 170 |
+
for r in results:
|
| 171 |
+
n = len(r["true"])
|
| 172 |
+
if n == 0: continue
|
| 173 |
+
n_cls = r["probs"].shape[1]
|
| 174 |
+
one_hot = np.zeros((n, n_cls))
|
| 175 |
+
one_hot[np.arange(n), r["true"]] = 1
|
| 176 |
+
scores.append(np.mean(np.sum((r["probs"] - one_hot) ** 2, axis=-1)))
|
| 177 |
+
return np.mean(scores) if scores else 0
|
| 178 |
+
|
| 179 |
+
# ===================== TEMP SCALING =====================
|
| 180 |
+
|
| 181 |
+
def find_temperature(cal_results, grid=None):
|
| 182 |
+
if grid is None:
|
| 183 |
+
grid = np.linspace(0.5, 5.0, 50)
|
| 184 |
+
all_logits, all_labels = [], []
|
| 185 |
+
for r in cal_results:
|
| 186 |
+
probs = np.clip(r["probs"], 1e-10, 1.0)
|
| 187 |
+
all_logits.append(np.log(probs))
|
| 188 |
+
all_labels.append(r["true"])
|
| 189 |
+
all_logits = np.concatenate(all_logits)
|
| 190 |
+
all_labels = np.concatenate(all_labels)
|
| 191 |
+
best_t, best_nll = 1.0, float("inf")
|
| 192 |
+
for t in grid:
|
| 193 |
+
scaled = all_logits / t
|
| 194 |
+
max_log = np.max(scaled, axis=-1, keepdims=True)
|
| 195 |
+
log_probs = scaled - max_log - np.log(np.sum(np.exp(scaled - max_log), axis=-1, keepdims=True))
|
| 196 |
+
nll = -np.mean(log_probs[np.arange(len(all_labels)), all_labels])
|
| 197 |
+
if nll < best_nll:
|
| 198 |
+
best_nll = nll
|
| 199 |
+
best_t = t
|
| 200 |
+
return best_t
|
| 201 |
+
|
| 202 |
+
def apply_temperature(results, temp):
|
| 203 |
+
scaled = []
|
| 204 |
+
for r in results:
|
| 205 |
+
probs = np.clip(r["probs"], 1e-10, 1.0)
|
| 206 |
+
logits = np.log(probs) / temp
|
| 207 |
+
max_log = np.max(logits, axis=-1, keepdims=True)
|
| 208 |
+
new_probs = np.exp(logits - max_log) / np.sum(np.exp(logits - max_log), axis=-1, keepdims=True)
|
| 209 |
+
scaled.append({"true": r["true"], "probs": new_probs, "preds": np.argmax(new_probs, axis=-1)})
|
| 210 |
+
return scaled
|
| 211 |
+
|
| 212 |
+
# ===================== CONFORMAL =====================
|
| 213 |
+
|
| 214 |
+
def conformal_qhat(cal_results, alpha=0.1):
|
| 215 |
+
scores = [1.0 - r["probs"][j, label] for r in cal_results for j, label in enumerate(r["true"])]
|
| 216 |
+
scores = np.array(scores)
|
| 217 |
+
n = len(scores)
|
| 218 |
+
q = np.ceil((n + 1) * (1 - alpha)) / n
|
| 219 |
+
return np.quantile(scores, q, method="higher")
|
| 220 |
+
|
| 221 |
+
def conformal_qhat_class_conditional(cal_results, alpha=0.1):
|
| 222 |
+
class_scores = defaultdict(list)
|
| 223 |
+
for r in cal_results:
|
| 224 |
+
for j, label in enumerate(r["true"]):
|
| 225 |
+
class_scores[label].append(1.0 - r["probs"][j, label])
|
| 226 |
+
thresholds = {}
|
| 227 |
+
for label, scores in class_scores.items():
|
| 228 |
+
scores = np.array(scores)
|
| 229 |
+
n = len(scores)
|
| 230 |
+
if n == 0:
|
| 231 |
+
thresholds[label] = 1.0
|
| 232 |
+
continue
|
| 233 |
+
q = np.ceil((n + 1) * (1 - alpha)) / n
|
| 234 |
+
thresholds[label] = np.quantile(scores, q, method="higher")
|
| 235 |
+
return thresholds
|
| 236 |
+
|
| 237 |
+
def evaluate_conformal(results, q_hat, n_classes, per_class_thresholds=None):
|
| 238 |
+
coverage_count, total = 0, 0
|
| 239 |
+
set_sizes = []
|
| 240 |
+
class_cov = defaultdict(int)
|
| 241 |
+
class_tot = defaultdict(int)
|
| 242 |
+
class_set = defaultdict(list)
|
| 243 |
+
size_strat = defaultdict(lambda: {"correct": 0, "total": 0})
|
| 244 |
+
for r in results:
|
| 245 |
+
for j, label in enumerate(r["true"]):
|
| 246 |
+
total += 1
|
| 247 |
+
threshold = per_class_thresholds.get(label, q_hat) if per_class_thresholds else q_hat
|
| 248 |
+
pred_set = [y for y in range(n_classes) if (1.0 - r["probs"][j, y]) <= threshold]
|
| 249 |
+
set_size = len(pred_set)
|
| 250 |
+
set_sizes.append(set_size)
|
| 251 |
+
size_strat[set_size]["total"] += 1
|
| 252 |
+
if label in pred_set:
|
| 253 |
+
coverage_count += 1
|
| 254 |
+
class_cov[label] += 1
|
| 255 |
+
size_strat[set_size]["correct"] += 1
|
| 256 |
+
class_tot[label] += 1
|
| 257 |
+
class_set[label].append(set_size)
|
| 258 |
+
coverage = coverage_count / total if total else 0
|
| 259 |
+
avg_size = np.mean(set_sizes) if set_sizes else 0
|
| 260 |
+
per_class = {}
|
| 261 |
+
for k in sorted(class_tot.keys()):
|
| 262 |
+
per_class[k] = {
|
| 263 |
+
"coverage": class_cov[k] / class_tot[k] if class_tot[k] else 0,
|
| 264 |
+
"avg_set_size": np.mean(class_set[k]) if class_set[k] else 0,
|
| 265 |
+
}
|
| 266 |
+
size_strat_out = {}
|
| 267 |
+
for size in sorted(size_strat.keys()):
|
| 268 |
+
d = size_strat[size]
|
| 269 |
+
size_strat_out[size] = {
|
| 270 |
+
"coverage": d["correct"] / d["total"] if d["total"] else 0,
|
| 271 |
+
"n": d["total"],
|
| 272 |
+
}
|
| 273 |
+
return coverage, avg_size, per_class, size_strat_out
|
| 274 |
+
|
| 275 |
+
def evaluate_mondrian(cal_results, test_results, alpha, n_classes):
|
| 276 |
+
class_cal = defaultdict(list)
|
| 277 |
+
for r in cal_results:
|
| 278 |
+
for j, label in enumerate(r["true"]):
|
| 279 |
+
class_cal[label].append(1.0 - r["probs"][j, label])
|
| 280 |
+
thresholds = {}
|
| 281 |
+
for label, scores in class_cal.items():
|
| 282 |
+
scores = np.array(scores)
|
| 283 |
+
n = len(scores)
|
| 284 |
+
if n == 0:
|
| 285 |
+
thresholds[label] = 1.0
|
| 286 |
+
continue
|
| 287 |
+
q = np.ceil((n + 1) * (1 - alpha)) / n
|
| 288 |
+
thresholds[label] = np.quantile(scores, q, method="higher")
|
| 289 |
+
class_cov = defaultdict(lambda: {"correct": 0, "total": 0})
|
| 290 |
+
class_set = defaultdict(list)
|
| 291 |
+
for r in test_results:
|
| 292 |
+
for j, label in enumerate(r["true"]):
|
| 293 |
+
threshold = thresholds.get(label, 1.0)
|
| 294 |
+
pred_set = [y for y in range(n_classes) if (1.0 - r["probs"][j, y]) <= threshold]
|
| 295 |
+
set_size = len(pred_set)
|
| 296 |
+
class_cov[label]["total"] += 1
|
| 297 |
+
class_set[label].append(set_size)
|
| 298 |
+
if label in pred_set:
|
| 299 |
+
class_cov[label]["correct"] += 1
|
| 300 |
+
mondrian = {}
|
| 301 |
+
for k in sorted(class_cov.keys()):
|
| 302 |
+
d = class_cov[k]
|
| 303 |
+
mondrian[k] = {
|
| 304 |
+
"coverage": d["correct"] / d["total"] if d["total"] else 0,
|
| 305 |
+
"avg_set_size": np.mean(class_set[k]) if class_set[k] else 0,
|
| 306 |
+
"n": d["total"],
|
| 307 |
+
}
|
| 308 |
+
return mondrian
|
| 309 |
+
|
| 310 |
+
# ===================== BASELINES =====================
|
| 311 |
+
|
| 312 |
+
def entropy_baseline(results, alpha, n_classes):
|
| 313 |
+
coverage_count, total = 0, 0
|
| 314 |
+
set_sizes = []
|
| 315 |
+
for r in results:
|
| 316 |
+
for j, label in enumerate(r["true"]):
|
| 317 |
+
total += 1
|
| 318 |
+
probs = r["probs"][j]
|
| 319 |
+
sorted_idx = np.argsort(-probs)
|
| 320 |
+
cumsum = np.cumsum(probs[sorted_idx])
|
| 321 |
+
n_include = np.searchsorted(cumsum, 1 - alpha) + 1
|
| 322 |
+
pred_set = sorted_idx[:n_include].tolist()
|
| 323 |
+
set_sizes.append(len(pred_set))
|
| 324 |
+
if label in pred_set:
|
| 325 |
+
coverage_count += 1
|
| 326 |
+
return coverage_count / total if total else 0, np.mean(set_sizes) if set_sizes else 0
|
| 327 |
+
|
| 328 |
+
def maxmargin_baseline(results, alpha, n_classes):
|
| 329 |
+
all_margins = []
|
| 330 |
+
for r in results:
|
| 331 |
+
for j in range(len(r["true"])):
|
| 332 |
+
probs = r["probs"][j]
|
| 333 |
+
sp = np.sort(probs)[::-1]
|
| 334 |
+
all_margins.append(sp[0] - sp[1] if len(sp) > 1 else 1.0)
|
| 335 |
+
all_margins = np.array(all_margins)
|
| 336 |
+
n = len(all_margins)
|
| 337 |
+
q = np.ceil((n + 1) * (1 - alpha)) / n
|
| 338 |
+
margin_thresh = np.quantile(all_margins, q, method="higher")
|
| 339 |
+
coverage_count, total = 0, 0
|
| 340 |
+
set_sizes = []
|
| 341 |
+
for r in results:
|
| 342 |
+
for j, label in enumerate(r["true"]):
|
| 343 |
+
total += 1
|
| 344 |
+
probs = r["probs"][j]
|
| 345 |
+
sorted_idx = np.argsort(-probs)
|
| 346 |
+
sp = np.sort(probs)[::-1]
|
| 347 |
+
margin = sp[0] - sp[1] if len(sp) > 1 else 1.0
|
| 348 |
+
if margin >= margin_thresh:
|
| 349 |
+
pred_set = [sorted_idx[0]]
|
| 350 |
+
else:
|
| 351 |
+
pred_set = sorted_idx[:min(2, n_classes)].tolist()
|
| 352 |
+
set_sizes.append(len(pred_set))
|
| 353 |
+
if label in pred_set:
|
| 354 |
+
coverage_count += 1
|
| 355 |
+
return coverage_count / total if total else 0, np.mean(set_sizes) if set_sizes else 0
|
| 356 |
+
|
| 357 |
+
# ===================== PRIORITIZATION =====================
|
| 358 |
+
|
| 359 |
+
def experiment_prioritization(results, budgets):
|
| 360 |
+
all_unc, all_errors = [], []
|
| 361 |
+
for r in results:
|
| 362 |
+
max_probs = np.max(r["probs"], axis=-1)
|
| 363 |
+
uncertainties = 1 - max_probs
|
| 364 |
+
errors = (r["preds"] != r["true"]).astype(float)
|
| 365 |
+
all_unc.extend(uncertainties)
|
| 366 |
+
all_errors.extend(errors)
|
| 367 |
+
all_unc = np.array(all_unc)
|
| 368 |
+
all_errors = np.array(all_errors)
|
| 369 |
+
n_total = len(all_unc)
|
| 370 |
+
out = {}
|
| 371 |
+
for budget in budgets:
|
| 372 |
+
b = min(budget, n_total)
|
| 373 |
+
random_idx = np.random.choice(n_total, size=b, replace=False)
|
| 374 |
+
random_rate = all_errors[random_idx].mean()
|
| 375 |
+
sorted_idx = np.argsort(-all_unc)
|
| 376 |
+
top_idx = sorted_idx[:b]
|
| 377 |
+
unc_rate = all_errors[top_idx].mean()
|
| 378 |
+
catch = unc_rate / random_rate if random_rate > 0 else float('inf')
|
| 379 |
+
out[budget] = {
|
| 380 |
+
"random_error_rate": float(random_rate),
|
| 381 |
+
"uncertainty_error_rate": float(unc_rate),
|
| 382 |
+
"catch_rate": float(catch),
|
| 383 |
+
}
|
| 384 |
+
return out
|
| 385 |
+
|
| 386 |
+
# ===================== PIPELINE =====================
|
| 387 |
+
|
| 388 |
+
def run_pipeline(model_id, dataset_loader, infer_fn, task_name, n_classes, label_map, budgets=[100, 500, 1000, 5000]):
|
| 389 |
+
log(f"\n{'='*60}")
|
| 390 |
+
log(f"TASK: {task_name}")
|
| 391 |
+
log(f"MODEL: {model_id}")
|
| 392 |
+
log(f"{'='*60}")
|
| 393 |
+
|
| 394 |
+
model, tokenizer = load_model(model_id)
|
| 395 |
+
cal_ds, test_ds = dataset_loader()
|
| 396 |
+
log(f" Calibration: {len(cal_ds)} seqs, Test: {len(test_ds)} seqs")
|
| 397 |
+
|
| 398 |
+
log(" Running inference (calibration)...")
|
| 399 |
+
cal_results = infer_fn(model, tokenizer, cal_ds)
|
| 400 |
+
log(f" Calibration residues: {sum(len(r['true']) for r in cal_results):,}")
|
| 401 |
+
|
| 402 |
+
log(" Running inference (test)...")
|
| 403 |
+
test_results = infer_fn(model, tokenizer, test_ds)
|
| 404 |
+
log(f" Test residues: {sum(len(r['true']) for r in test_results):,}")
|
| 405 |
+
|
| 406 |
+
del model
|
| 407 |
+
|
| 408 |
+
# Baseline
|
| 409 |
+
base_acc = compute_accuracy(test_results)
|
| 410 |
+
base_ece = compute_ece(test_results)
|
| 411 |
+
base_brier = compute_brier(test_results)
|
| 412 |
+
log(f" Baseline: Acc={base_acc:.4f}, ECE={base_ece:.4f}, Brier={base_brier:.4f}")
|
| 413 |
+
|
| 414 |
+
# Temperature scaling
|
| 415 |
+
best_t = find_temperature(cal_results)
|
| 416 |
+
scaled_cal = apply_temperature(cal_results, best_t)
|
| 417 |
+
scaled_test = apply_temperature(test_results, best_t)
|
| 418 |
+
ts_acc = compute_accuracy(scaled_test)
|
| 419 |
+
ts_ece = compute_ece(scaled_test)
|
| 420 |
+
ts_brier = compute_brier(scaled_test)
|
| 421 |
+
ece_red = (base_ece - ts_ece) / base_ece * 100 if base_ece else 0
|
| 422 |
+
log(f" Temperature T={best_t:.2f}: Acc={ts_acc:.4f}, ECE={ts_ece:.4f} ({ece_red:+.0f}%), Brier={ts_brier:.4f}")
|
| 423 |
+
|
| 424 |
+
# Conformal (raw)
|
| 425 |
+
log(" Conformal prediction...")
|
| 426 |
+
conformal = {}
|
| 427 |
+
for alpha in [0.05, 0.10, 0.20]:
|
| 428 |
+
q = conformal_qhat(cal_results, alpha)
|
| 429 |
+
cov, size, pclass, sstrat = evaluate_conformal(test_results, q, n_classes)
|
| 430 |
+
log(f" Raw alpha={alpha:.2f}: cov={cov:.4f}, set={size:.2f}")
|
| 431 |
+
q_s = conformal_qhat(scaled_cal, alpha)
|
| 432 |
+
cov_s, size_s, pclass_s, sstrat_s = evaluate_conformal(scaled_test, q_s, n_classes)
|
| 433 |
+
log(f" T-scaled alpha={alpha:.2f}: cov={cov_s:.4f}, set={size_s:.2f}")
|
| 434 |
+
conformal[f"alpha_{alpha}"] = {
|
| 435 |
+
"raw": {"coverage": float(cov), "avg_set_size": float(size),
|
| 436 |
+
"per_class": {label_map.get(k, str(k)): v for k, v in pclass.items()},
|
| 437 |
+
"size_stratified": {str(kk): vv for kk, vv in sstrat.items()}},
|
| 438 |
+
"temperature_scaled": {"coverage": float(cov_s), "avg_set_size": float(size_s),
|
| 439 |
+
"per_class": {label_map.get(k, str(k)): v for k, v in pclass_s.items()},
|
| 440 |
+
"size_stratified": {str(kk): vv for kk, vv in sstrat_s.items()}},
|
| 441 |
+
}
|
| 442 |
+
|
| 443 |
+
# Class-conditional
|
| 444 |
+
log(" Class-conditional conformal...")
|
| 445 |
+
cc = {}
|
| 446 |
+
for alpha in [0.05, 0.10, 0.20]:
|
| 447 |
+
th = conformal_qhat_class_conditional(cal_results, alpha)
|
| 448 |
+
cov, size, pclass, _ = evaluate_conformal(test_results, 0, n_classes, th)
|
| 449 |
+
log(f" alpha={alpha:.2f}: cov={cov:.4f}, set={size:.2f}")
|
| 450 |
+
cc[f"alpha_{alpha}"] = {
|
| 451 |
+
"coverage": float(cov), "avg_set_size": float(size),
|
| 452 |
+
"per_class": {label_map.get(k, str(k)): v for k, v in pclass.items()},
|
| 453 |
+
}
|
| 454 |
+
|
| 455 |
+
# Mondrian
|
| 456 |
+
log(" Mondrian conformal...")
|
| 457 |
+
mondrian = {}
|
| 458 |
+
for alpha in [0.05, 0.10, 0.20]:
|
| 459 |
+
mon = evaluate_mondrian(cal_results, test_results, alpha, n_classes)
|
| 460 |
+
log(f" alpha={alpha:.2f}")
|
| 461 |
+
for k, v in mon.items():
|
| 462 |
+
log(f" {label_map.get(k, str(k))}: cov={v['coverage']:.4f}, set={v['avg_set_size']:.2f}, n={v['n']}")
|
| 463 |
+
mondrian[f"alpha_{alpha}"] = {label_map.get(k, str(k)): v for k, v in mon.items()}
|
| 464 |
+
|
| 465 |
+
# Baselines
|
| 466 |
+
log(" Baselines...")
|
| 467 |
+
ent = {}
|
| 468 |
+
mm = {}
|
| 469 |
+
for alpha in [0.05, 0.10, 0.20]:
|
| 470 |
+
ec, es = entropy_baseline(test_results, alpha, n_classes)
|
| 471 |
+
mc, ms = maxmargin_baseline(test_results, alpha, n_classes)
|
| 472 |
+
log(f" alpha={alpha:.2f}: Entropy cov={ec:.4f} set={es:.2f}, MaxMargin cov={mc:.4f} set={ms:.2f}")
|
| 473 |
+
ent[f"alpha_{alpha}"] = {"coverage": float(ec), "avg_set_size": float(es)}
|
| 474 |
+
mm[f"alpha_{alpha}"] = {"coverage": float(mc), "avg_set_size": float(ms)}
|
| 475 |
+
|
| 476 |
+
# Prioritization
|
| 477 |
+
log(" Experiment prioritization...")
|
| 478 |
+
prio = experiment_prioritization(test_results, budgets)
|
| 479 |
+
for b, d in prio.items():
|
| 480 |
+
log(f" Budget={b}: random={d['random_error_rate']:.3f}, unc={d['uncertainty_error_rate']:.3f}, catch={d['catch_rate']:.2f}x")
|
| 481 |
+
|
| 482 |
+
return {
|
| 483 |
+
"task": task_name,
|
| 484 |
+
"model": model_id,
|
| 485 |
+
"baseline": {"accuracy": float(base_acc), "ece": float(base_ece), "brier": float(base_brier)},
|
| 486 |
+
"temperature_scaling": {"temperature": float(best_t), "accuracy": float(ts_acc),
|
| 487 |
+
"ece": float(ts_ece), "brier": float(ts_brier),
|
| 488 |
+
"ece_reduction_pct": float(ece_red)},
|
| 489 |
+
"conformal": conformal,
|
| 490 |
+
"class_conditional": cc,
|
| 491 |
+
"mondrian": mondrian,
|
| 492 |
+
"entropy_baseline": ent,
|
| 493 |
+
"maxmargin_baseline": mm,
|
| 494 |
+
"experiment_prioritization": prio,
|
| 495 |
+
"_cal_raw": cal_results,
|
| 496 |
+
"_cal_scaled": scaled_cal,
|
| 497 |
+
"_test_raw": test_results,
|
| 498 |
+
"_test_scaled": scaled_test,
|
| 499 |
+
}
|
| 500 |
+
|
| 501 |
+
# ===================== MAIN =====================
|
| 502 |
+
|
| 503 |
+
def main():
|
| 504 |
+
log("=" * 60)
|
| 505 |
+
log("ConformalESM Job 1: Core Experiments (8M SS + 35M Disorder)")
|
| 506 |
+
log("CPU-only, all post-hoc, no retraining")
|
| 507 |
+
log("=" * 60)
|
| 508 |
+
|
| 509 |
+
all_results = {}
|
| 510 |
+
|
| 511 |
+
# Task 1: Secondary Structure - 8M
|
| 512 |
+
ss8m = run_pipeline(SS_MODEL, load_ss_data, infer_ss,
|
| 513 |
+
"Secondary Structure (Q3) - ESM-2-8M", 3, SS_ID2LABEL)
|
| 514 |
+
all_results["ss_8m"] = {k: v for k, v in ss8m.items() if not k.startswith("_")}
|
| 515 |
+
|
| 516 |
+
# Task 2: Disorder - 35M
|
| 517 |
+
dis35m = run_pipeline(DIS_MODEL, load_disorder_data, infer_disorder,
|
| 518 |
+
"Disorder Prediction - ESM-2-35M", 2, {0: "Ordered", 1: "Disordered"})
|
| 519 |
+
all_results["disorder_35m"] = {k: v for k, v in dis35m.items() if not k.startswith("_")}
|
| 520 |
+
|
| 521 |
+
# Save cal/test raw results for cross-model transfer in Job 2
|
| 522 |
+
raw_data = {
|
| 523 |
+
"ss_8m_cal_raw": ss8m["_cal_raw"],
|
| 524 |
+
"ss_8m_cal_scaled": ss8m["_cal_scaled"],
|
| 525 |
+
"ss_8m_test_raw": ss8m["_test_raw"],
|
| 526 |
+
"ss_8m_test_scaled": ss8m["_test_scaled"],
|
| 527 |
+
"dis_35m_cal_raw": dis35m["_cal_raw"],
|
| 528 |
+
"dis_35m_cal_scaled": dis35m["_cal_scaled"],
|
| 529 |
+
"dis_35m_test_raw": dis35m["_test_raw"],
|
| 530 |
+
"dis_35m_test_scaled": dis35m["_test_scaled"],
|
| 531 |
+
}
|
| 532 |
+
|
| 533 |
+
# Save results
|
| 534 |
+
log(f"\n{'='*60}")
|
| 535 |
+
log("Saving Results")
|
| 536 |
+
log(f"{'='*60}")
|
| 537 |
+
|
| 538 |
+
with open("job1_results.json", "w") as f:
|
| 539 |
+
json.dump(all_results, f, indent=2)
|
| 540 |
+
log(" Saved: job1_results.json")
|
| 541 |
+
|
| 542 |
+
# Push to hub
|
| 543 |
+
log(" Pushing to knoxel/conformalesm-paper-starter...")
|
| 544 |
+
try:
|
| 545 |
+
from huggingface_hub import HfApi
|
| 546 |
+
api = HfApi()
|
| 547 |
+
api.upload_file(
|
| 548 |
+
path_or_fileobj="job1_results.json",
|
| 549 |
+
path_in_repo="job1_results.json",
|
| 550 |
+
repo_id="knoxel/conformalesm-paper-starter",
|
| 551 |
+
repo_type="model",
|
| 552 |
+
)
|
| 553 |
+
log(" Successfully pushed Job 1 results to Hub!")
|
| 554 |
+
except Exception as e:
|
| 555 |
+
log(f" Could not push to Hub: {e}")
|
| 556 |
+
|
| 557 |
+
log(f"\n{'='*60}")
|
| 558 |
+
log("JOB 1 COMPLETE")
|
| 559 |
+
log(f"{'='*60}")
|
| 560 |
+
|
| 561 |
+
|
| 562 |
+
if __name__ == "__main__":
|
| 563 |
+
main()
|