Upload phaseA_new_methods.py with huggingface_hub

phaseA_new_methods.py

@@ -0,0 +1,222 @@
+"""
+New mapping methods on top of the existing 25-anchor experiment data:
+
+  (f) procrustes              Per-tensor orthogonal Procrustes alignment between
+                              X-anchor and Y-anchor matrices in tensor-flat space.
+                              Respects the LoRA gauge ambiguity better than ridge.
+  (g) topk_global_ridge       Pick top-K anchors by X-side cosine similarity to X_target,
+                              then run global anchor-basis ridge on that subset.
+  (h) topk_pertensor_ridge    Same, but per-tensor ridge.
+  (i) topk_pertensor_mlp      Top-K + per-tensor PCA-MLP hypernet (re-trained per task).
+
+Reuses all training artifacts from /app/scaled/X and /app/scaled/Y.
+"""
+import os, json, gc, shutil, re, collections, math
+from pathlib import Path
+import numpy as np
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
+from peft import PeftModel
+from safetensors.torch import load_file, save_file
+
+# Reuse main pipeline's task definitions
+import sys; sys.path.insert(0, "/app")
+import scaled_pipeline as sp
+
+set_seed(42)
+
+OUT = sp.OUT  # /app/scaled
+ANCHOR_NAMES = sp.ANCHOR_NAMES
+HELDOUT_NAMES = sp.HELDOUT_NAMES
+MODEL_Y = sp.MODEL_Y
+EVAL_PER_TASK = sp.EVAL_PER_TASK
+
+def load_sd(p): return {k: v.float().cpu() for k,v in load_file(str(p/"adapter_model.safetensors")).items()}
+
+def parse_key(k):
+    m = re.search(r'layers\.(\d+)\.self_attn\.(\w+)_proj\.lora_(A|B)', k)
+    return int(m.group(1)), m.group(2), m.group(3)
+
+def group_by_layer(sd):
+    g = collections.defaultdict(dict)
+    for k, v in sd.items():
+        L, mod, AB = parse_key(k)
+        g[(mod, AB)][L] = (k, v)
+    return g
+
+def align_X_to_Y(X_anchors_list, X_target, Y_template_sd):
+    Y_g = group_by_layer(Y_template_sd)
+    X_gs = [group_by_layer(s) for s in X_anchors_list]
+    X_t  = group_by_layer(X_target)
+    nLY = max(Y_g[("q","A")].keys()) + 1
+    nLX = max(X_gs[0][("q","A")].keys()) + 1
+    out = {}
+    for (mod, AB), layers in Y_g.items():
+        for L_y, (yk, _) in layers.items():
+            L_x = round(L_y * (nLX-1) / max(1, nLY-1))
+            x_anchors = [X_gs[i][(mod,AB)][L_x][1] for i in range(len(X_gs))]
+            x_target  = X_t[(mod,AB)][L_x][1]
+            out[yk] = (x_anchors, x_target)
+    return out
+
+# ----------------- (f) Procrustes -----------------
+def procrustes_pred(X_anchors, Y_anchors, X_target):
+    """Per-tensor orthogonal Procrustes (rectangular), avoiding dx*dy materialization.
+    Cross-covariance M = Xc^T Yc has rank ≤ N. Compute its SVD efficiently:
+        Xc = U_x S_x V_x^T  (compact, V_x^T is (N, dx), V_x is (dx, N))
+        Yc = U_y S_y V_y^T
+        Then SVD(M) <- via N-by-N inner SVD; never materialize dx-by-dy.
+    """
+    align = align_X_to_Y(X_anchors, X_target, Y_anchors[0])
+    out = {}
+    for yk, (x_anchors, x_target) in align.items():
+        Y_mat = torch.stack([s[yk].reshape(-1) for s in Y_anchors])  # (N, dy)
+        X_mat = torch.stack([t.reshape(-1) for t in x_anchors])      # (N, dx)
+        Ym, Xm = Y_mat.mean(0), X_mat.mean(0)
+        Yc, Xc = Y_mat - Ym, X_mat - Xm
+        Ux_, Sx_, Vxh = torch.linalg.svd(Xc, full_matrices=False)
+        Uy_, Sy_, Vyh = torch.linalg.svd(Yc, full_matrices=False)
+        C = (Sx_.unsqueeze(1) * Ux_.T) @ (Uy_ * Sy_.unsqueeze(0))   # (N, N)
+        Uz, Sz, Vzh = torch.linalg.svd(C, full_matrices=False)
+        scale = (Sz.sum() / (Sx_.pow(2).sum() + 1e-8)).item()
+        xt_c = x_target.reshape(-1) - Xm
+        a = Vxh @ xt_c                  # (N,)
+        b = (Uz @ Vzh) @ a              # (N,)
+        delta = b @ Vyh                 # (dy,)
+        out[yk] = (Ym + scale * delta).reshape(Y_anchors[0][yk].shape)
+    return out
+
+# ----------------- (g/h/i) top-K nearest anchor selection -----------------
+def topk_anchor_idx(X_anchors, X_target, k=8):
+    """Return indices of top-K anchors by cosine similarity to X_target."""
+    keys = sorted(X_anchors[0].keys())
+    Xa = torch.stack([torch.cat([s[kk].flatten() for kk in keys]) for s in X_anchors])  # (N, dx)
+    xt = torch.cat([X_target[kk].flatten() for kk in keys])
+    sims = F.cosine_similarity(Xa, xt.unsqueeze(0), dim=1)
+    idx = torch.topk(sims, k=min(k, len(X_anchors))).indices.tolist()
+    return idx, sims.tolist()
+
+def topk_global_ridge(X_anchors, Y_anchors, X_target, k=8, ridge=1e-3):
+    idx, sims = topk_anchor_idx(X_anchors, X_target, k=k)
+    Xa = [X_anchors[i] for i in idx]
+    Ya = [Y_anchors[i] for i in idx]
+    pred, alpha = sp.global_ridge_pred(Xa, Ya, X_target, ridge=ridge)
+    return pred, idx, sims
+
+def topk_pertensor_ridge(X_anchors, Y_anchors, X_target, k=8, ridge=1e-3):
+    idx, sims = topk_anchor_idx(X_anchors, X_target, k=k)
+    Xa = [X_anchors[i] for i in idx]
+    Ya = [Y_anchors[i] for i in idx]
+    pred = sp.pertensor_ridge_pred(Xa, Ya, X_target, ridge=ridge)
+    return pred, idx, sims
+
+def topk_pertensor_mlp(X_anchors, Y_anchors, X_target, k=12, k_lat=6, epochs=300, lr=1e-3):
+    idx, sims = topk_anchor_idx(X_anchors, X_target, k=k)
+    Xa = [X_anchors[i] for i in idx]
+    Ya = [Y_anchors[i] for i in idx]
+    pred, _ = sp.pertensor_pca_mlp_pred(Xa, Ya, X_target, k_lat=k_lat, epochs=epochs, lr=lr)
+    return pred, idx, sims
+
+# ----------------- save / eval -----------------
+def save_adapter(sd, dirname, src_template):
+    d = OUT/"Y_pred"/dirname
+    d.mkdir(parents=True, exist_ok=True)
+    shutil.copy(src_template/"adapter_config.json", d/"adapter_config.json")
+    save_file({k: v.to(torch.bfloat16) for k,v in sd.items()}, str(d/"adapter_model.safetensors"))
+    return d
+
+@torch.no_grad()
+def eval_adapter(adapter_dir, eval_ds, labels, tok, max_n=300):
+    base = AutoModelForCausalLM.from_pretrained(MODEL_Y, torch_dtype=torch.bfloat16, attn_implementation="eager").cuda()
+    if adapter_dir is None:
+        m = base
+    else:
+        m = PeftModel.from_pretrained(base, str(adapter_dir))
+    acc = sp.eval_classification(m, tok, eval_ds, labels, max_n=max_n)
+    del m, base
+    if adapter_dir is None: pass
+    gc.collect(); torch.cuda.empty_cache()
+    return acc
+
+def cos_sd(sd1, sd2):
+    keys = sorted(sd1.keys())
+    a = torch.cat([sd1[k].float().flatten() for k in keys])
+    b = torch.cat([sd2[k].float().flatten() for k in keys])
+    return F.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item()
+
+# ----------------- main -----------------
+def main():
+    print("Loading anchors...")
+    X_anchors = [load_sd(OUT/"X"/n) for n in ANCHOR_NAMES]
+    Y_anchors = [load_sd(OUT/"Y"/n) for n in ANCHOR_NAMES]
+    print(f"  N={len(X_anchors)} anchors")
+
+    tokY = AutoTokenizer.from_pretrained(MODEL_Y)
+    if tokY.pad_token is None: tokY.pad_token = tokY.eos_token
+    tokY.padding_side = "left"
+
+    out_results = {}
+    template = OUT/"Y"/ANCHOR_NAMES[0]
+
+    for t_name in HELDOUT_NAMES:
+        print(f"\n=== Held-out task: {t_name} ===")
+        X_target = load_sd(OUT/"X"/t_name)
+        Y_oracle = load_sd(OUT/"Y"/t_name)
+
+        results = {"cosines": {}, "acc": {}, "selected_anchors": {}}
+
+        # build predictions
+        preds = {}
+        print("  procrustes ...")
+        preds["procrustes"] = procrustes_pred(X_anchors, Y_anchors, X_target)
+
+        for k_top in [5, 8, 12]:
+            print(f"  topk_global_ridge k={k_top} ...")
+            p, idx, sims = topk_global_ridge(X_anchors, Y_anchors, X_target, k=k_top)
+            preds[f"topk{k_top}_global_ridge"] = p
+            results["selected_anchors"][f"topk{k_top}"] = [ANCHOR_NAMES[i] for i in idx]
+            print(f"    selected: {[ANCHOR_NAMES[i] for i in idx]}")
+
+            print(f"  topk_pertensor_ridge k={k_top} ...")
+            p, _, _ = topk_pertensor_ridge(X_anchors, Y_anchors, X_target, k=k_top)
+            preds[f"topk{k_top}_pertensor_ridge"] = p
+
+        print("  topk12_pertensor_mlp ...")
+        p, _, _ = topk_pertensor_mlp(X_anchors, Y_anchors, X_target, k=12, k_lat=6, epochs=300)
+        preds["topk12_pertensor_mlp"] = p
+
+        # save and eval
+        _, eval_ds, labels = sp.build_task(t_name, n_train=10, n_eval=EVAL_PER_TASK)
+
+        # cosines
+        for n, sd in preds.items():
+            results["cosines"][n] = cos_sd(sd, Y_oracle)
+
+        # accuracies
+        results["acc"]["base_Y"] = eval_adapter(None, eval_ds, labels, tokY)
+        print(f"  base_Y={results['acc']['base_Y']:.4f}")
+        for n, sd in preds.items():
+            d = save_adapter(sd, f"{t_name}_{n}", template)
+            acc = eval_adapter(d, eval_ds, labels, tokY)
+            results["acc"][n] = acc
+            print(f"  {n:32s}  cos={results['cosines'][n]:.4f}  acc={acc:.4f}")
+
+        results["acc"]["oracle_Y"] = eval_adapter(OUT/"Y"/t_name, eval_ds, labels, tokY)
+        print(f"  oracle_Y={results['acc']['oracle_Y']:.4f}")
+
+        out_results[t_name] = results
+
+    # aggregate
+    methods = ["base_Y","procrustes","topk5_global_ridge","topk5_pertensor_ridge",
+               "topk8_global_ridge","topk8_pertensor_ridge",
+               "topk12_global_ridge","topk12_pertensor_ridge","topk12_pertensor_mlp","oracle_Y"]
+    avg = {m: float(np.mean([out_results[t]["acc"][m] for t in HELDOUT_NAMES])) for m in methods}
+    print("\n=== AVG (new methods) ===")
+    for m in methods: print(f"  {m:32s}  {avg[m]:.4f}")
+    out_results["avg"] = avg
+    (OUT/"results_phaseA.json").write_text(json.dumps(out_results, indent=2, default=float))
+
+if __name__ == "__main__":
+    main()