knoxel commited on
Commit
22f1aab
·
verified ·
1 Parent(s): 79405d9

Upload run_conformalesm_job1.py

Browse files
Files changed (1) hide show
  1. 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()