Buckets:
| #!/usr/bin/env python3 | |
| """Three-way comparison: SAE-50 vs SAE-150 vs Raw-hidden. Response tokens only, identical scoring. | |
| Reports per-topic + overall F1(=Dice) / IoU / Accuracy(recall), and how much of the | |
| SAE-50 -> Raw gap the 150-feature budget closes. | |
| """ | |
| from __future__ import annotations | |
| import json, sys | |
| from pathlib import Path | |
| import numpy as np, torch | |
| sys.path.insert(0, "/workspace/scope/modeling") | |
| from models import PAD, TOPICS, build_model, topic_overlap | |
| from sklearn.metrics import f1_score, classification_report, accuracy_score | |
| ROOT = Path("/workspace/scope") | |
| SAE500 = ROOT / "cache/sae500_2k_clean_prompt_response" | |
| SAE1500 = ROOT / "cache/sae1500_2k_clean_prompt_response" | |
| RAW = ROOT / "cache/hidden10_2k_prefill" | |
| RESP = 2 | |
| ARMS = { # family -> {arm: (ckpt, cache)} | |
| "GRU": { | |
| "50": (ROOT/"results/gru_2k/gru_lr3e4/checkpoint_best.pt", SAE500), | |
| "150": (ROOT/"results/sae1500_2k/gru/gru_lr3e4/checkpoint_best.pt", SAE1500), | |
| "raw": (ROOT/"ablation/gru_hidden/gru_hidden_no_sae/checkpoint_best.pt", RAW)}, | |
| "Transformer": { | |
| "50": (ROOT/"results/transformer_2k/transformer_wide/checkpoint_best.pt", SAE500), | |
| "150": (ROOT/"results/sae1500_2k/transformer/transformer_wide/checkpoint_best.pt", SAE1500), | |
| "raw": (ROOT/"ablation/baselines_hidden_transformer/transformer_wide/checkpoint_best.pt", RAW)}, | |
| "ConvNeXt": { | |
| "50": (ROOT/"results/tcn_2k_clean/convnext/checkpoint.pt", SAE500), | |
| "150": (ROOT/"results/sae1500_2k/tcn/convnext/checkpoint_best.pt", SAE1500), | |
| "raw": (ROOT/"ablation/baselines_hidden_convnext/convnext/checkpoint_best.pt", RAW)}, | |
| } | |
| def load(ckpt, device): | |
| p = torch.load(ckpt, map_location="cpu", weights_only=False) | |
| m = build_model(p["architecture"], p["model"]); m.load_state_dict(p["state_dict"]) | |
| return m.to(device).eval(), sum(x.numel() for x in m.parameters()) | |
| def score(ckpt, cache, device): | |
| m, params = load(ckpt, device) | |
| feats = np.load(cache/"features.npy", mmap_mode="r") | |
| labels = np.load(cache/"labels.npy"); roles = np.load(cache/"role_ids.npy") | |
| lengths = np.load(cache/"lengths.npy"); off = np.r_[0, np.cumsum(lengths)] | |
| test = np.load(cache/"split_indices.npz")["test"] | |
| T, P = [], [] | |
| with torch.inference_mode(): | |
| for i in test: | |
| i = int(i); lo, hi = int(off[i]), int(off[i+1]) | |
| x = torch.from_numpy(np.asarray(feats[lo:hi], dtype=np.float32))[None].to(device) | |
| pred = m(x)[0].argmax(-1).cpu().numpy() | |
| lab = labels[lo:hi].astype(np.int64).copy(); lab[roles[lo:hi] != RESP] = PAD | |
| keep = lab != PAD | |
| T.append(lab[keep]); P.append(pred[keep]) | |
| truth = np.concatenate(T); predk = np.concatenate(P) | |
| ov = topic_overlap(truth, predk) | |
| rep = classification_report(truth, predk, labels=range(len(TOPICS)), | |
| target_names=TOPICS, output_dict=True, zero_division=0) | |
| per = {} | |
| for t in TOPICS: | |
| tm = truth == TOPICS.index(t) | |
| per[t] = {"f1": rep[t]["f1-score"], "iou": ov["per_topic"][t]["iou"], | |
| "acc": float((predk[tm] == TOPICS.index(t)).mean()) if tm.any() else 0.0} | |
| return {"params": params, "acc": float(accuracy_score(truth, predk)), | |
| "macro_f1": float(f1_score(truth, predk, average="macro", zero_division=0)), | |
| "macro_iou": ov["macro_iou"], "per_topic": per} | |
| def gap(o, n, r): | |
| g = r - o | |
| return 100 * (n - o) / g if abs(g) > 1e-9 else float("nan") | |
| def main(): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| out = {} | |
| for fam, arms in ARMS.items(): | |
| out[fam] = {a: score(c, cache, device) for a, (c, cache) in arms.items()} | |
| print(f"scored {fam}", flush=True) | |
| json.dump(out, open(ROOT/"results/sae1500_2k/three_way.json", "w"), indent=2) | |
| fams = ["GRU", "Transformer", "ConvNeXt"] | |
| print("\n==================== OVERALL: 50 / 150 / raw (150 closes X% of 50->raw gap) ====================") | |
| for label, key in [("Accuracy", "acc"), ("Macro F1 (=Dice)", "macro_f1"), ("Macro IoU", "macro_iou")]: | |
| print(f"\n -- {label} --") | |
| print(f" {'model':12s} {'SAE-50':>8s} {'SAE-150':>8s} {'RAW':>8s} {'150 vs raw':>11s} {'%gap closed':>11s}") | |
| for f in fams: | |
| a, n, r = out[f]["50"][key], out[f]["150"][key], out[f]["raw"][key] | |
| print(f" {f:12s} {a:8.4f} {n:8.4f} {r:8.4f} {n-r:+11.4f} {gap(a,n,r):9.0f}%") | |
| print(f"\n -- Parameters (M) --") | |
| print(f" {'model':12s} {'SAE-50':>8s} {'SAE-150':>8s} {'RAW':>8s} {'raw/150':>8s}") | |
| for f in fams: | |
| a, n, r = out[f]["50"]["params"]/1e6, out[f]["150"]["params"]/1e6, out[f]["raw"]["params"]/1e6 | |
| print(f" {f:12s} {a:8.2f} {n:8.2f} {r:8.2f} {r/n:7.1f}x") | |
| for label, key in [("F1 (=Dice)", "f1"), ("IoU", "iou"), ("Accuracy", "acc")]: | |
| print(f"\n==================== PER-TOPIC {label}: 50 / 150 / raw (%gap closed) ====================") | |
| print(f"{'topic':22s}" + "".join(f"{f[:10]:>26s}" for f in fams)) | |
| print(f"{'':22s}" + "".join(f"{'50 150 raw %gap':>26s}" for f in fams)) | |
| for t in TOPICS: | |
| line = f"{t:22s}" | |
| for f in fams: | |
| a = out[f]["50"]["per_topic"][t][key]; n = out[f]["150"]["per_topic"][t][key]; r = out[f]["raw"]["per_topic"][t][key] | |
| line += f" {a:.3f} {n:.3f} {r:.3f} {gap(a,n,r):4.0f}%" | |
| print(line) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 5.44 kB
- Xet hash:
- b236b0d09aa011715aa03354a49f60491536bf130d9e70d492ba9946e5bdca97
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.