phaseA_new_methods.py

"""
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()