Upload pipeline.py with huggingface_hub

pipeline.py

@@ -0,0 +1,342 @@
+"""
+Cross-Model LoRA Adapter Prediction
+====================================
+Model X = Qwen/Qwen2.5-0.5B-Instruct
+Model Y = meta-llama/Llama-3.2-1B-Instruct
+Tasks   : A=SST-2, B=AG News, C=Subj, D=Emotion (held out for Y)
+
+Pipeline:
+  1. Train LoRA X_A, X_B, X_C, X_D on Model X
+  2. Train LoRA Y_A, Y_B, Y_C, Y_D on Model Y     (Y_D = oracle, kept aside)
+  3. Learn mapping f from {X_A,X_B,X_C} -> {Y_A,Y_B,Y_C} via anchor-basis ridge regression
+  4. Predict Y_hat_D = f(X_D)
+  5. Evaluate on D test split: base Y, mean(Y_A,Y_B,Y_C), Y_hat_D (predicted), Y_D (oracle), Y_D trained-on-X-train (sanity)
+  6. Push everything to HF repo
+"""
+import os, json, gc, math, time, argparse, shutil
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from datasets import load_dataset, Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
+from peft import LoraConfig, get_peft_model, PeftModel
+from peft.utils import get_peft_model_state_dict, set_peft_model_state_dict
+from trl import SFTTrainer, SFTConfig
+
+set_seed(42)
+
+# -------------------- Config --------------------
+MODEL_X = "Qwen/Qwen2.5-0.5B-Instruct"
+MODEL_Y = "meta-llama/Llama-3.2-1B-Instruct"
+
+LORA_R = 8
+LORA_ALPHA = 16
+LORA_TARGETS = ["q_proj", "v_proj"]
+
+TRAIN_PER_TASK = 1500   # SFT examples per task
+EVAL_PER_TASK = 400     # eval examples
+EPOCHS = 1
+BS = 8
+LR = 2e-4
+MAX_LEN = 192
+
+OUT = Path("/app/out")
+OUT.mkdir(exist_ok=True, parents=True)
+
+# -------------------- Datasets --------------------
+def fmt(prompt: str, label_text: str):
+    return [
+        {"role": "user", "content": prompt},
+        {"role": "assistant", "content": label_text},
+    ]
+
+def build_task(name: str):
+    """Return (train_ds, eval_ds, label_set, prompt_fn) where each row has a 'messages' field."""
+    if name == "A":  # SST-2 sentiment
+        ds = load_dataset("stanfordnlp/sst2")
+        labels = ["negative", "positive"]
+        def to_msg(r): return {"messages": fmt(
+            f"Classify the sentiment of this sentence as 'negative' or 'positive'. Respond with just the label.\n\nSentence: {r['sentence'].strip()}\n\nSentiment:",
+            labels[r["label"]])}
+        train = ds["train"].shuffle(seed=0).select(range(TRAIN_PER_TASK)).map(to_msg, remove_columns=ds["train"].column_names)
+        ev = ds["validation"].shuffle(seed=0).select(range(min(EVAL_PER_TASK, len(ds["validation"])))).map(to_msg, remove_columns=ds["validation"].column_names)
+        return train, ev, labels, "sentiment"
+    if name == "B":  # AG News
+        ds = load_dataset("fancyzhx/ag_news")
+        labels = ["world", "sports", "business", "sci/tech"]
+        def to_msg(r): return {"messages": fmt(
+            f"Classify the news topic as 'world', 'sports', 'business', or 'sci/tech'. Respond with just the label.\n\nNews: {r['text'].strip()}\n\nTopic:",
+            labels[r["label"]])}
+        train = ds["train"].shuffle(seed=0).select(range(TRAIN_PER_TASK)).map(to_msg, remove_columns=ds["train"].column_names)
+        ev = ds["test"].shuffle(seed=0).select(range(EVAL_PER_TASK)).map(to_msg, remove_columns=ds["test"].column_names)
+        return train, ev, labels, "topic"
+    if name == "C":  # Subj
+        ds = load_dataset("SetFit/subj")
+        labels = ["objective", "subjective"]
+        def to_msg(r): return {"messages": fmt(
+            f"Classify whether this sentence is 'objective' or 'subjective'. Respond with just the label.\n\nSentence: {r['text'].strip()}\n\nLabel:",
+            labels[r["label"]])}
+        train = ds["train"].shuffle(seed=0).select(range(min(TRAIN_PER_TASK, len(ds["train"])))).map(to_msg, remove_columns=ds["train"].column_names)
+        ev = ds["test"].shuffle(seed=0).select(range(min(EVAL_PER_TASK, len(ds["test"])))).map(to_msg, remove_columns=ds["test"].column_names)
+        return train, ev, labels, "subjectivity"
+    if name == "D":  # Emotion
+        ds = load_dataset("dair-ai/emotion", "split")
+        labels = ["sadness", "joy", "love", "anger", "fear", "surprise"]
+        def to_msg(r): return {"messages": fmt(
+            f"Classify the emotion of this sentence as one of: 'sadness', 'joy', 'love', 'anger', 'fear', 'surprise'. Respond with just the label.\n\nSentence: {r['text'].strip()}\n\nEmotion:",
+            labels[r["label"]])}
+        train = ds["train"].shuffle(seed=0).select(range(TRAIN_PER_TASK)).map(to_msg, remove_columns=ds["train"].column_names)
+        ev = ds["test"].shuffle(seed=0).select(range(EVAL_PER_TASK)).map(to_msg, remove_columns=ds["test"].column_names)
+        return train, ev, labels, "emotion"
+    raise ValueError(name)
+
+TASKS = ["A", "B", "C", "D"]
+
+# -------------------- Train one LoRA --------------------
+def train_lora(model_name: str, task: str, save_dir: Path):
+    if save_dir.exists() and (save_dir/"adapter_model.safetensors").exists():
+        print(f"[SKIP] {save_dir} already exists")
+        return
+    save_dir.mkdir(parents=True, exist_ok=True)
+    print(f"\n=== Training LoRA: model={model_name} task={task} -> {save_dir}")
+    tok = AutoTokenizer.from_pretrained(model_name)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, attn_implementation="eager")
+    model.config.use_cache = False
+    train_ds, _, _, _ = build_task(task)
+    lora = LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, target_modules=LORA_TARGETS,
+                      lora_dropout=0.0, bias="none", task_type="CAUSAL_LM")
+    cfg = SFTConfig(
+        output_dir=str(save_dir/"_trainer"),
+        num_train_epochs=EPOCHS,
+        per_device_train_batch_size=BS,
+        gradient_accumulation_steps=1,
+        learning_rate=LR,
+        lr_scheduler_type="cosine",
+        warmup_ratio=0.05,
+        bf16=True,
+        max_seq_length=MAX_LEN,
+        logging_steps=25,
+        logging_first_step=True,
+        logging_strategy="steps",
+        disable_tqdm=True,
+        save_strategy="no",
+        report_to="none",
+        seed=42,
+        packing=False,
+    )
+    trainer = SFTTrainer(model=model, args=cfg, train_dataset=train_ds, peft_config=lora, tokenizer=tok)
+    trainer.train()
+    trainer.model.save_pretrained(str(save_dir))
+    tok.save_pretrained(str(save_dir))
+    # cleanup
+    shutil.rmtree(save_dir/"_trainer", ignore_errors=True)
+    del trainer, model
+    gc.collect(); torch.cuda.empty_cache()
+
+# -------------------- Cross-model mapping --------------------
+def load_adapter_state(path: Path) -> Dict[str, torch.Tensor]:
+    """Load LoRA state dict, kept on CPU as float32."""
+    from safetensors.torch import load_file
+    sd = load_file(str(path/"adapter_model.safetensors"))
+    return {k: v.float().cpu() for k, v in sd.items()}
+
+def flatten_sd(sd: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, List[Tuple[str, torch.Size]]]:
+    keys_shapes = [(k, sd[k].shape) for k in sorted(sd.keys())]
+    flat = torch.cat([sd[k].reshape(-1) for k, _ in keys_shapes])
+    return flat, keys_shapes
+
+def unflatten(flat: torch.Tensor, keys_shapes) -> Dict[str, torch.Tensor]:
+    out = {}
+    i = 0
+    for k, shape in keys_shapes:
+        n = int(np.prod(shape))
+        out[k] = flat[i:i+n].reshape(shape)
+        i += n
+    return out
+
+def predict_anchor_basis(X_anchors: List[torch.Tensor], Y_anchors: List[torch.Tensor],
+                         X_target: torch.Tensor, ridge: float = 1e-3) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    f maps X-side -> Y-side using a paired-anchor linear basis.
+
+    Express x_target - mean(X) ≈ sum_i alpha_i (X_i - mean(X)) via ridge regression
+    on the small 3x3 Gram matrix; then ŷ = mean(Y) + sum_i alpha_i (Y_i - mean(Y)).
+
+    Returns (y_hat, alpha).
+    """
+    Xs = torch.stack(X_anchors)         # [k, dx]
+    Ys = torch.stack(Y_anchors)         # [k, dy]
+    Xm = Xs.mean(0); Ym = Ys.mean(0)
+    Xc = Xs - Xm                        # [k, dx]
+    Yc = Ys - Ym                        # [k, dy]
+    xc = X_target - Xm                  # [dx]
+    G = Xc @ Xc.T                       # [k, k]
+    rhs = Xc @ xc                       # [k]
+    alpha = torch.linalg.solve(G + ridge * torch.eye(G.shape[0]), rhs)  # [k]
+    y_hat = Ym + (alpha @ Yc)           # [dy]
+    return y_hat, alpha
+
+# -------------------- Evaluation --------------------
+@torch.no_grad()
+def eval_classification(model, tok, eval_ds, labels: List[str], max_new=8) -> float:
+    """Greedy generation; compare first non-empty token-stripped substring against labels."""
+    model.eval()
+    correct = 0; total = 0
+    label_set = [l.lower() for l in labels]
+    bs = 16
+    prompts = []
+    golds = []
+    for ex in eval_ds:
+        msgs = ex["messages"]
+        gold = msgs[1]["content"].strip().lower()
+        # build prompt up to assistant turn
+        prompt = tok.apply_chat_template([msgs[0]], tokenize=False, add_generation_prompt=True)
+        prompts.append(prompt)
+        golds.append(gold)
+    for i in range(0, len(prompts), bs):
+        batch = prompts[i:i+bs]
+        enc = tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=MAX_LEN).to(model.device)
+        out = model.generate(**enc, max_new_tokens=max_new, do_sample=False, pad_token_id=tok.pad_token_id)
+        gen = out[:, enc["input_ids"].shape[1]:]
+        for j, g in enumerate(gen):
+            txt = tok.decode(g, skip_special_tokens=True).strip().lower()
+            # match longest prefix label
+            pred = None
+            for lab in sorted(label_set, key=len, reverse=True):
+                if txt.startswith(lab):
+                    pred = lab; break
+            if pred is None:
+                # fallback: any label appearing
+                for lab in label_set:
+                    if lab in txt: pred = lab; break
+            if pred == golds[i+j]:
+                correct += 1
+            total += 1
+    return correct / max(1,total)
+
+
+# -------------------- Main --------------------
+def main(stage: str = "all"):
+    # ---------- Stage 1+2: train all LoRAs ----------
+    if stage in ("all", "train"):
+        for t in TASKS:
+            train_lora(MODEL_X, t, OUT/"X"/f"X_{t}")
+        for t in TASKS:  # train Y_D too for oracle
+            train_lora(MODEL_Y, t, OUT/"Y"/f"Y_{t}")
+
+    # ---------- Stage 3: build mapping + predict Y_hat_D ----------
+    if stage in ("all", "map"):
+        print("\n=== Building cross-model mapping ===")
+        X_states = {t: load_adapter_state(OUT/"X"/f"X_{t}") for t in TASKS}
+        Y_states = {t: load_adapter_state(OUT/"Y"/f"Y_{t}") for t in TASKS}
+
+        # flatten — same keys/shapes within each side
+        X_flat = {}; Y_flat = {}
+        Xks = Yks = None
+        for t in TASKS:
+            f, ks = flatten_sd(X_states[t]);  X_flat[t] = f; Xks = ks
+            f, ks = flatten_sd(Y_states[t]);  Y_flat[t] = f; Yks = ks
+        print("X adapter dim:", X_flat["A"].numel(), "Y adapter dim:", Y_flat["A"].numel())
+
+        # anchor-basis ridge regression mapping
+        Xa = [X_flat["A"], X_flat["B"], X_flat["C"]]
+        Ya = [Y_flat["A"], Y_flat["B"], Y_flat["C"]]
+        Y_hat_D, alpha = predict_anchor_basis(Xa, Ya, X_flat["D"], ridge=1e-3)
+        print("Anchor weights alpha (A,B,C):", alpha.tolist())
+        # also: mean baseline
+        Y_mean_ABC = torch.stack(Ya).mean(0)
+        # cosine sim diagnostics
+        def cos(a, b): return torch.nn.functional.cosine_similarity(a.flatten().unsqueeze(0), b.flatten().unsqueeze(0)).item()
+        print("cos(Y_hat_D, Y_D) =", cos(Y_hat_D, Y_flat["D"]))
+        print("cos(Y_mean_ABC, Y_D) =", cos(Y_mean_ABC, Y_flat["D"]))
+        print("cos(Y_A, Y_D) =", cos(Y_flat["A"], Y_flat["D"]))
+        print("cos(Y_B, Y_D) =", cos(Y_flat["B"], Y_flat["D"]))
+        print("cos(Y_C, Y_D) =", cos(Y_flat["C"], Y_flat["D"]))
+
+        # save predicted + mean adapters as standard PEFT checkpoints (clone Y_A's metadata)
+        from safetensors.torch import save_file
+        for name, flat in [("Y_pred_D", Y_hat_D), ("Y_mean_ABC", Y_mean_ABC)]:
+            sd = unflatten(flat, Yks)
+            sd_bf16 = {k: v.to(torch.bfloat16) for k, v in sd.items()}
+            d = OUT/"Y"/name
+            d.mkdir(parents=True, exist_ok=True)
+            # copy adapter_config and tokenizer from Y_A
+            shutil.copy(OUT/"Y"/"Y_A"/"adapter_config.json", d/"adapter_config.json")
+            for f in ["tokenizer.json","tokenizer_config.json","special_tokens_map.json"]:
+                src = OUT/"Y"/"Y_A"/f
+                if src.exists(): shutil.copy(src, d/f)
+            save_file(sd_bf16, str(d/"adapter_model.safetensors"))
+            print("Saved", d)
+
+        # save mapping diagnostics
+        diag = {
+            "alpha_ABC": alpha.tolist(),
+            "cos_Yhat_YD": cos(Y_hat_D, Y_flat["D"]),
+            "cos_Ymean_YD": cos(Y_mean_ABC, Y_flat["D"]),
+            "cos_YA_YD": cos(Y_flat["A"], Y_flat["D"]),
+            "cos_YB_YD": cos(Y_flat["B"], Y_flat["D"]),
+            "cos_YC_YD": cos(Y_flat["C"], Y_flat["D"]),
+            "X_dim": X_flat["A"].numel(),
+            "Y_dim": Y_flat["A"].numel(),
+            "ridge": 1e-3,
+        }
+        (OUT/"mapping_diagnostics.json").write_text(json.dumps(diag, indent=2))
+
+    # ---------- Stage 4: evaluate on D ----------
+    if stage in ("all", "eval"):
+        print("\n=== Evaluating on task D (Emotion) ===")
+        _, eval_d, labels_d, _ = build_task("D")
+        tok = AutoTokenizer.from_pretrained(MODEL_Y)
+        if tok.pad_token is None: tok.pad_token = tok.eos_token
+        tok.padding_side = "left"
+        results = {}
+        # Base Y
+        base = AutoModelForCausalLM.from_pretrained(MODEL_Y, torch_dtype=torch.bfloat16, attn_implementation="eager").cuda()
+        results["base_Y"] = eval_classification(base, tok, eval_d, labels_d)
+        print("base_Y", results["base_Y"])
+        del base; gc.collect(); torch.cuda.empty_cache()
+
+        # helper for adapter eval
+        def with_adapter(adapter_dir):
+            base = AutoModelForCausalLM.from_pretrained(MODEL_Y, torch_dtype=torch.bfloat16, attn_implementation="eager").cuda()
+            m = PeftModel.from_pretrained(base, str(adapter_dir))
+            acc = eval_classification(m, tok, eval_d, labels_d)
+            del m, base; gc.collect(); torch.cuda.empty_cache()
+            return acc
+
+        for name, dirname in [
+            ("Y_A_on_D", "Y_A"),
+            ("Y_B_on_D", "Y_B"),
+            ("Y_C_on_D", "Y_C"),
+            ("Y_mean_ABC_on_D", "Y_mean_ABC"),
+            ("Y_pred_D", "Y_pred_D"),
+            ("Y_oracle_D", "Y_D"),
+        ]:
+            results[name] = with_adapter(OUT/"Y"/dirname)
+            print(name, results[name])
+
+        # also: sanity-check Model X with X_D oracle on its own dataset
+        tokx = AutoTokenizer.from_pretrained(MODEL_X)
+        if tokx.pad_token is None: tokx.pad_token = tokx.eos_token
+        tokx.padding_side = "left"
+        basex = AutoModelForCausalLM.from_pretrained(MODEL_X, torch_dtype=torch.bfloat16, attn_implementation="eager").cuda()
+        results["base_X"] = eval_classification(basex, tokx, eval_d, labels_d)
+        del basex; gc.collect(); torch.cuda.empty_cache()
+        basex = AutoModelForCausalLM.from_pretrained(MODEL_X, torch_dtype=torch.bfloat16, attn_implementation="eager").cuda()
+        mx = PeftModel.from_pretrained(basex, str(OUT/"X"/"X_D"))
+        results["X_oracle_D"] = eval_classification(mx, tokx, eval_d, labels_d)
+        del mx, basex; gc.collect(); torch.cuda.empty_cache()
+
+        (OUT/"results.json").write_text(json.dumps(results, indent=2))
+        print("\n=== Results ===")
+        for k, v in results.items(): print(f"  {k:24s} {v:.4f}")
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--stage", default="all", choices=["all","train","map","eval"])
+    args = ap.parse_args()
+    main(args.stage)