Upload reasoning pipeline script

reasoning_pipeline.py

@@ -1 +1,438 @@
-# empty placeholder; real upload below
+"""
+Cross-Model LoRA Adapter Prediction — Reasoning validation (GSM8K held-out).
+
+Pipeline:
+  1) Build 9 reasoning datasets in a uniform "Q -> rationale + #### answer" format.
+  2) Train LoRA adapters on Qwen2.5-0.5B-Instruct (X) and Llama-3.2-1B-Instruct (Y) for
+     8 anchor tasks + GSM8K (oracle) per model. Y for GSM8K is held-out from the mapping.
+  3) Fit mappings (mean, global_ridge, topk_global_ridge K=4) on the 8 anchor pairs and
+     predict Y_GSM8K from X_GSM8K.
+  4) Eval base_Y, mean, global_ridge, topk_global_ridge, oracle_Y on 250 GSM8K test problems
+     using greedy generation + numeric-answer match on the final number after '####'.
+  5) Save adapters + results.json under /app/reasoning/, then push to the Hub.
+"""
+
+import os, re, json, gc, time, math, argparse, random, shutil
+from pathlib import Path
+from typing import List, Dict, Tuple
+
+import numpy as np
+import torch
+from datasets import load_dataset, Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import LoraConfig, PeftModel, get_peft_model
+from trl import SFTTrainer, SFTConfig
+
+ROOT = Path("/app/reasoning")
+ROOT.mkdir(parents=True, exist_ok=True)
+
+MODEL_X = "Qwen/Qwen2.5-0.5B-Instruct"
+MODEL_Y = "meta-llama/Llama-3.2-1B-Instruct"
+
+LORA_KWARGS = dict(r=8, lora_alpha=16, lora_dropout=0.0, target_modules=["q_proj", "v_proj"], bias="none")
+
+# ---------------- dataset preparation ----------------
+def _final_num(s: str):
+    nums = re.findall(r"-?\d+\.?\d*", s.replace(",", ""))
+    return nums[-1] if nums else None
+
+def fmt(question: str, answer_text: str, gold: str) -> Dict:
+    """Return chat-format messages. assistant ends with '#### gold'."""
+    rationale = answer_text.strip()
+    if "####" not in rationale:
+        rationale = f"{rationale}\n#### {gold}"
+    return {
+        "messages": [
+            {"role": "user", "content": question.strip()},
+            {"role": "assistant", "content": rationale},
+        ],
+        "gold": str(gold),
+    }
+
+def load_gsm8k(split, n):
+    ds = load_dataset("openai/gsm8k", "main", split=split)
+    ds = ds.shuffle(seed=0).select(range(min(n, len(ds))))
+    out = []
+    for r in ds:
+        gold = _final_num(r["answer"].split("####")[-1])
+        out.append(fmt(r["question"], r["answer"], gold))
+    return out
+
+def load_svamp(n):
+    ds = load_dataset("ChilleD/SVAMP", split="train").shuffle(seed=0).select(range(min(n, 700)))
+    out = []
+    for r in ds:
+        q = (r["Body"] + " " + r["Question"]).strip()
+        gold = str(r["Answer"]).rstrip(".0") if "." in str(r["Answer"]) else str(r["Answer"])
+        try:
+            f = float(r["Answer"])
+            gold = str(int(f)) if f.is_integer() else str(f)
+        except Exception:
+            gold = str(r["Answer"])
+        rationale = f"{r['Equation']} = {gold}"
+        out.append(fmt(q, rationale, gold))
+    return out
+
+def load_multiarith(n):
+    ds = load_dataset("ChilleD/MultiArith", split="train").shuffle(seed=0).select(range(min(n, 400)))
+    out = []
+    for r in ds:
+        gold = str(r["final_ans"]).strip()
+        out.append(fmt(r["question"], f"The answer is {gold}.", gold))
+    return out
+
+def load_aqua(n):
+    ds = load_dataset("deepmind/aqua_rat", "raw", split="train").shuffle(seed=0).select(range(min(n, 800)))
+    out = []
+    for r in ds:
+        opts = "\n".join(r["options"])
+        q = f"{r['question']}\nOptions:\n{opts}"
+        gold = str(r["correct"]).strip()
+        out.append(fmt(q, r["rationale"], gold))
+    return out
+
+def load_mathplus(n):
+    ds = load_dataset("TIGER-Lab/MATH-plus", split="train").shuffle(seed=0).select(range(min(n*3, 3000)))
+    out = []
+    for r in ds:
+        out_text = r["output"]
+        # extract a final number/expression — fallback: last token
+        gold = _final_num(out_text) or out_text.strip().split()[-1][:32]
+        out.append(fmt(r["instruction"], out_text, gold))
+        if len(out) >= n:
+            break
+    return out
+
+def load_strategyqa(n):
+    ds = load_dataset("tasksource/strategy-qa", split="train").shuffle(seed=0).select(range(min(n, 2000)))
+    out = []
+    for r in ds:
+        gold = "yes" if r["answer"] else "no"
+        facts = " ".join(r["facts"])
+        rationale = facts
+        out.append(fmt(r["question"] + "\nAnswer yes or no.", rationale, gold))
+    return out
+
+def _mc_format(stem, choices, label):
+    if isinstance(choices, dict):
+        texts, labels = choices["text"], choices["label"]
+    else:
+        texts, labels = choices.text, choices.label
+    body = "\n".join(f"{l}. {t}" for l, t in zip(labels, texts))
+    return f"{stem}\n{body}\nAnswer with the letter only.", str(label).strip()
+
+def load_openbookqa(n):
+    ds = load_dataset("allenai/openbookqa", "main", split="train").shuffle(seed=0).select(range(min(n, 4000)))
+    out = []
+    for r in ds:
+        q, gold = _mc_format(r["question_stem"], r["choices"], r["answerKey"])
+        out.append(fmt(q, f"The answer is {gold}.", gold))
+    return out
+
+def load_arc(n, cfg):
+    ds = load_dataset("allenai/ai2_arc", cfg, split="train").shuffle(seed=0).select(range(min(n, 2000)))
+    out = []
+    for r in ds:
+        q, gold = _mc_format(r["question"], r["choices"], r["answerKey"])
+        out.append(fmt(q, f"The answer is {gold}.", gold))
+    return out
+
+DATASETS = {
+    "svamp":       lambda n: load_svamp(n),
+    "multiarith":  lambda n: load_multiarith(n),
+    "aqua":        lambda n: load_aqua(n),
+    "mathplus":    lambda n: load_mathplus(n),
+    "strategyqa":  lambda n: load_strategyqa(n),
+    "openbookqa":  lambda n: load_openbookqa(n),
+    "arc_easy":    lambda n: load_arc(n, "ARC-Easy"),
+    "arc_chall":   lambda n: load_arc(n, "ARC-Challenge"),
+    "gsm8k":       lambda n: load_gsm8k("train", n),
+}
+
+ANCHOR_TASKS = ["svamp", "multiarith", "aqua", "mathplus", "strategyqa", "openbookqa", "arc_easy", "arc_chall"]
+HELD_OUT = "gsm8k"
+
+# ---------------- training ----------------
+def train_one(model_id: str, task: str, examples: List[Dict], out_dir: Path, max_len: int = 384, n_train: int = 800):
+    out_dir.mkdir(parents=True, exist_ok=True)
+    if (out_dir / "adapter_config.json").exists():
+        print(f"  [skip] {out_dir} already exists")
+        return
+    tok = AutoTokenizer.from_pretrained(model_id)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="cuda")
+    ds = Dataset.from_list([{"messages": e["messages"]} for e in examples[:n_train]])
+    cfg = SFTConfig(
+        output_dir=str(out_dir / "_trainer"),
+        per_device_train_batch_size=8,
+        gradient_accumulation_steps=1,
+        num_train_epochs=1,
+        learning_rate=2e-4,
+        lr_scheduler_type="cosine",
+        warmup_ratio=0.05,
+        bf16=True,
+        max_seq_length=max_len,
+        logging_strategy="steps",
+        logging_steps=20,
+        logging_first_step=True,
+        save_strategy="no",
+        report_to="none",
+        disable_tqdm=True,
+        gradient_checkpointing=False,
+    )
+    peft_cfg = LoraConfig(task_type="CAUSAL_LM", **LORA_KWARGS)
+    trainer = SFTTrainer(
+        model=model,
+        args=cfg,
+        train_dataset=ds,
+        peft_config=peft_cfg,
+        processing_class=tok,
+    )
+    trainer.train()
+    trainer.model.save_pretrained(str(out_dir))
+    tok.save_pretrained(str(out_dir))
+    shutil.rmtree(out_dir / "_trainer", ignore_errors=True)
+    del trainer, model
+    gc.collect(); torch.cuda.empty_cache()
+
+# ---------------- adapter I/O ----------------
+def load_adapter_state(path: Path) -> Dict[str, torch.Tensor]:
+    """Return {param_name: tensor (cpu, float32)} for the LoRA-only weights."""
+    from safetensors.torch import load_file
+    f = path / "adapter_model.safetensors"
+    sd = load_file(str(f))
+    return {k: v.detach().to(torch.float32).cpu() for k, v in sd.items()}
+
+def save_adapter_state(reference_dir: Path, new_state: Dict[str, torch.Tensor], out_dir: Path):
+    """Copy adapter_config.json from reference, write new safetensors."""
+    from safetensors.torch import save_file
+    out_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copy(reference_dir / "adapter_config.json", out_dir / "adapter_config.json")
+    for extra in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]:
+        if (reference_dir / extra).exists():
+            shutil.copy(reference_dir / extra, out_dir / extra)
+    save_file({k: v.to(torch.bfloat16).contiguous() for k, v in new_state.items()},
+              str(out_dir / "adapter_model.safetensors"))
+
+def flatten(state: Dict[str, torch.Tensor]) -> Tuple[np.ndarray, List[Tuple[str, tuple]]]:
+    keys = sorted(state.keys())
+    vecs, schema = [], []
+    for k in keys:
+        t = state[k]
+        schema.append((k, tuple(t.shape)))
+        vecs.append(t.flatten().numpy())
+    return np.concatenate(vecs), schema
+
+def unflatten(vec: np.ndarray, schema) -> Dict[str, torch.Tensor]:
+    out, off = {}, 0
+    for k, shp in schema:
+        n = int(np.prod(shp))
+        out[k] = torch.from_numpy(vec[off:off+n].reshape(shp).copy())
+        off += n
+    return out
+
+# ---------------- mappings ----------------
+def mean_baseline(Y_anchors: np.ndarray, X_target: np.ndarray) -> np.ndarray:
+    return Y_anchors.mean(axis=0)
+
+def global_ridge(X_anchors: np.ndarray, Y_anchors: np.ndarray, X_target: np.ndarray, lam=1e-3) -> np.ndarray:
+    Xb, Yb = X_anchors.mean(0), Y_anchors.mean(0)
+    Xc, Yc = X_anchors - Xb, Y_anchors - Yb
+    N = Xc.shape[0]
+    G = Xc @ Xc.T  # (N,N)
+    alpha = np.linalg.solve(G + lam * np.eye(N), Xc @ (X_target - Xb))
+    return Yb + alpha @ Yc
+
+def topk_global_ridge(X_anchors, Y_anchors, X_target, K=4, lam=1e-3):
+    Xb_full = X_anchors.mean(0)
+    sims = []
+    for i in range(X_anchors.shape[0]):
+        a = X_anchors[i] - Xb_full
+        b = X_target - Xb_full
+        denom = (np.linalg.norm(a) * np.linalg.norm(b)) + 1e-12
+        sims.append(float(a @ b) / denom)
+    sims = np.array(sims)
+    idx = np.argsort(-sims)[:K]
+    return global_ridge(X_anchors[idx], Y_anchors[idx], X_target, lam=lam), idx.tolist()
+
+# ---------------- eval ----------------
+def extract_pred(text: str, mode: str):
+    """mode = 'num' or 'letter' or 'yesno'."""
+    if mode == "num":
+        if "####" in text:
+            text = text.split("####")[-1]
+        return _final_num(text)
+    if mode == "letter":
+        m = re.search(r"\b([A-E])\b", text.upper())
+        return m.group(1) if m else None
+    if mode == "yesno":
+        t = text.lower()
+        if "yes" in t and "no" not in t.split("yes")[0][-10:]:
+            return "yes"
+        if "no" in t:
+            return "no"
+        return None
+    return text.strip()
+
+def eval_adapter(base_id: str, adapter_dir: Path | None, eval_set: List[Dict], mode: str, max_new=200):
+    tok = AutoTokenizer.from_pretrained(base_id, padding_side="left")
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=torch.bfloat16, device_map="cuda")
+    if adapter_dir is not None:
+        model = PeftModel.from_pretrained(model, str(adapter_dir))
+    model.eval()
+    correct = 0
+    preds = []
+    bs = 8
+    for i in range(0, len(eval_set), bs):
+        batch = eval_set[i:i+bs]
+        prompts = [tok.apply_chat_template([e["messages"][0]], tokenize=False, add_generation_prompt=True) for e in batch]
+        enc = tok(prompts, return_tensors="pt", padding=True, truncation=True, max_length=512).to("cuda")
+        with torch.no_grad():
+            out = model.generate(**enc, max_new_tokens=max_new, do_sample=False, pad_token_id=tok.pad_token_id)
+        gen = tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
+        for e, g in zip(batch, gen):
+            p = extract_pred(g, mode)
+            gold = e["gold"]
+            if mode == "num":
+                try:
+                    ok = (p is not None) and (abs(float(p) - float(gold)) < 1e-3)
+                except Exception:
+                    ok = False
+            else:
+                ok = (p is not None) and (str(p).strip().lower() == str(gold).strip().lower())
+            correct += int(ok)
+            preds.append({"gold": gold, "pred": p, "raw": g[:200]})
+    acc = correct / len(eval_set)
+    del model; gc.collect(); torch.cuda.empty_cache()
+    return acc, preds
+
+# ---------------- orchestrator ----------------
+def stage_train(n_train=800, n_per_task=900):
+    for task in ANCHOR_TASKS + [HELD_OUT]:
+        print(f"\n=== {task} ===")
+        examples = DATASETS[task](n_per_task)
+        print(f"  examples: {len(examples)}")
+        for tag, mid in [("X", MODEL_X), ("Y", MODEL_Y)]:
+            outd = ROOT / tag / task
+            print(f"  training {tag}/{task} -> {outd}")
+            t0 = time.time()
+            train_one(mid, task, examples, outd, n_train=n_train)
+            print(f"  done in {time.time()-t0:.1f}s")
+
+def stage_predict():
+    # load all anchor X, Y
+    X_states = {t: load_adapter_state(ROOT / "X" / t) for t in ANCHOR_TASKS + [HELD_OUT]}
+    Y_states = {t: load_adapter_state(ROOT / "Y" / t) for t in ANCHOR_TASKS + [HELD_OUT]}
+    # use Y schema for output
+    Y_ref_dir = ROOT / "Y" / ANCHOR_TASKS[0]
+    Y_vecs, Y_schema = {}, None
+    for t, s in Y_states.items():
+        v, sch = flatten(s)
+        Y_vecs[t] = v
+        Y_schema = sch
+    X_vecs, _ = {}, None
+    for t, s in X_states.items():
+        v, _ = flatten(s)
+        X_vecs[t] = v
+    Xa = np.stack([X_vecs[t] for t in ANCHOR_TASKS])
+    Ya = np.stack([Y_vecs[t] for t in ANCHOR_TASKS])
+    Xt = X_vecs[HELD_OUT]
+    Yt_oracle = Y_vecs[HELD_OUT]
+
+    preds = {}
+    preds["mean"] = mean_baseline(Ya, Xt)
+    preds["global_ridge"] = global_ridge(Xa, Ya, Xt)
+    yhat, idx = topk_global_ridge(Xa, Ya, Xt, K=4)
+    preds["topk4_global_ridge"] = yhat
+    print("topk4 indices:", [ANCHOR_TASKS[i] for i in idx])
+
+    cos = lambda a, b: float(a @ b / (np.linalg.norm(a)*np.linalg.norm(b)+1e-12))
+    cosines = {m: cos(v, Yt_oracle) for m, v in preds.items()}
+    print("cosines:", cosines)
+
+    out_pred = ROOT / "Y_pred"
+    for m, v in preds.items():
+        save_adapter_state(Y_ref_dir, unflatten(v, Y_schema), out_pred / m)
+    return cosines
+
+def stage_eval(n_eval=250):
+    eval_set = load_gsm8k("test", n_eval)
+    print(f"eval set size: {len(eval_set)}")
+    results = {}
+    # base Y (no adapter)
+    print("evaluating base Y ...")
+    acc, _ = eval_adapter(MODEL_Y, None, eval_set, mode="num")
+    results["base_Y"] = acc; print("  base_Y:", acc)
+    # base X
+    print("evaluating base X ...")
+    acc, _ = eval_adapter(MODEL_X, None, eval_set, mode="num")
+    results["base_X"] = acc; print("  base_X:", acc)
+    # oracle X
+    print("evaluating oracle X ...")
+    acc, _ = eval_adapter(MODEL_X, ROOT / "X" / HELD_OUT, eval_set, mode="num")
+    results["oracle_X"] = acc; print("  oracle_X:", acc)
+    # oracle Y
+    print("evaluating oracle Y ...")
+    acc, _ = eval_adapter(MODEL_Y, ROOT / "Y" / HELD_OUT, eval_set, mode="num")
+    results["oracle_Y"] = acc; print("  oracle_Y:", acc)
+    # predicted Y
+    for m in ["mean", "global_ridge", "topk4_global_ridge"]:
+        d = ROOT / "Y_pred" / m
+        if not (d / "adapter_config.json").exists():
+            continue
+        print(f"evaluating predicted Y/{m} ...")
+        acc, _ = eval_adapter(MODEL_Y, d, eval_set, mode="num")
+        results[m] = acc; print(f"  {m}:", acc)
+    return results
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--stage", default="all", choices=["train", "predict", "eval", "all", "smoke"])
+    ap.add_argument("--n_train", type=int, default=800)
+    ap.add_argument("--n_eval", type=int, default=250)
+    ap.add_argument("--push_repo", type=str, default=None)
+    args = ap.parse_args()
+
+    random.seed(0); np.random.seed(0); torch.manual_seed(0)
+
+    if args.stage == "smoke":
+        # tiny: 2 anchors, 50 examples, 20 eval
+        ANCHOR_TASKS[:] = ["svamp", "multiarith"]
+        stage_train(n_train=50, n_per_task=60)
+        cos = stage_predict()
+        res = stage_eval(n_eval=20)
+        print("SMOKE results:", res, "cos:", cos)
+    else:
+        all_results = {}
+        if args.stage in ("train", "all"):
+            stage_train(n_train=args.n_train)
+        if args.stage in ("predict", "all"):
+            all_results["cosines"] = stage_predict()
+        if args.stage in ("eval", "all"):
+            all_results["accuracy"] = stage_eval(n_eval=args.n_eval)
+        (ROOT / "results.json").write_text(json.dumps(all_results, indent=2))
+        print("\nFINAL:", json.dumps(all_results, indent=2))
+
+        if args.push_repo:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            api.create_repo(args.push_repo, exist_ok=True)
+            print(f"Uploading reasoning/ to {args.push_repo} ...")
+            api.upload_folder(
+                folder_path=str(ROOT),
+                path_in_repo="reasoning",
+                repo_id=args.push_repo,
+                repo_type="model",
+                commit_message="Add cross-model LoRA adapter prediction — reasoning validation (GSM8K)",
+            )
+            api.upload_file(
+                path_or_fileobj="/app/reasoning_pipeline.py",
+                path_in_repo="reasoning_pipeline.py",
+                repo_id=args.push_repo,
+                repo_type="model",
+                commit_message="Add reasoning pipeline script",
+            )
+            print("Push complete.")