J-lens workspace reproduction + loop retrofit: lens, band looping, adapters, controls, multi-task evals

Reproduction of the 2026 workspace/J-lens paper on gemma-4 (E2B/12B/26B), plus the workspace-loop retrofit line: merge adapter, prompt-only latent planning (MBPP), carry variant, attribution controls (FF/pause/untrained), band-location ablation, Blocksworld harness, 12B replication scripts. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-14 00:54:12 +02:00
commit ef9c08966c
42 changed files with 5068 additions and 0 deletions
@@ -0,0 +1,132 @@
+"""STaR-style difficulty labeling of MBPP with the frozen base model.
+
+Two passes per item: direct code generation vs plan-first-then-code, each
+executed against MBPP's unit tests (sandboxed subprocess). Labels:
+easy (direct passes), hard (plan-only passes), drop (neither). The model's
+own passing code is kept as the training target (in-distribution supervision).
+"""
+
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import torch
+from datasets import load_dataset
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+OUT.mkdir(exist_ok=True)
+
+DIRECT_SUFFIX = ("\n\nWrite only the Python function in a ```python code "
+                 "block. No explanation.")
+PLAN_SUFFIX = ("\n\nFirst write a very brief plan: at most 4 short bullet "
+               "lines, no headings, no math notation. Then write the complete "
+               "Python function in a ```python code block.")
+
+CODE_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.S)
+
+
+def mbpp_prompt(tok, item, suffix):
+    tests = "\n".join(item["test_list"])
+    msg = (f"{item['text']}\nYour code should pass these tests:\n\n{tests}"
+           f"{suffix}")
+    return tok.apply_chat_template([{"role": "user", "content": msg}],
+                                   tokenize=False, add_generation_prompt=True)
+
+
+def extract_code(text):
+    m = CODE_RE.findall(text)
+    return m[-1].strip() if m else None
+
+
+def run_tests(code, item, timeout=10):
+    if not code:
+        return False
+    script = (item.get("test_setup_code") or "") + "\n" + code + "\n" + \
+        "\n".join(item["test_list"])
+    try:
+        with tempfile.TemporaryDirectory() as td:
+            r = subprocess.run([sys.executable, "-c", script], cwd=td,
+                               capture_output=True, timeout=timeout)
+        return r.returncode == 0
+    except (subprocess.TimeoutExpired, OSError):
+        return False
+
+
+@torch.no_grad()
+def batch_generate(model, tok, prompts, max_new_tokens, batch_size=24):
+    outs = []
+    for i in range(0, len(prompts), batch_size):
+        chunk = prompts[i : i + batch_size]
+        enc = tok(chunk, return_tensors="pt", padding=True,
+                  add_special_tokens=False).to("cuda")
+        gen = model.generate(**enc, max_new_tokens=max_new_tokens,
+                             do_sample=False,
+                             pad_token_id=tok.pad_token_id or 0)
+        for j in range(len(chunk)):
+            outs.append(tok.decode(gen[j, enc["input_ids"].shape[1]:],
+                                   skip_special_tokens=True))
+        print(f"  {min(i+batch_size, len(prompts))}/{len(prompts)}", flush=True)
+    return outs
+
+
+def main():
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+    ds = load_dataset("google-research-datasets/mbpp", "full")
+
+    items = []
+    for split, tag in (("train", "train"), ("validation", "train"),
+                       ("test", "test")):
+        for row in ds[split]:
+            items.append({"split": tag, "task_id": row["task_id"],
+                          "text": row["text"], "test_list": row["test_list"],
+                          "test_setup_code": row["test_setup_code"]})
+    print(f"{len(items)} items", flush=True)
+
+    for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 220),
+                            ("plan", PLAN_SUFFIX, 700)):
+        t0 = time.time()
+        print(f"{tag} pass", flush=True)
+        gens = batch_generate(model, tok,
+                              [mbpp_prompt(tok, it, suffix) for it in items], mx)
+        codes = [extract_code(g) for g in gens]
+        with ThreadPoolExecutor(8) as ex:
+            oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
+                              zip(codes, items)))
+        for it, c, ok in zip(items, codes, oks):
+            it[f"{tag}_ok"] = bool(ok)
+            it[f"{tag}_code"] = c if ok else None
+        print(f"{tag} pass done in {time.time()-t0:.0f}s  "
+              f"pass@1={sum(oks)/len(items):.3f}", flush=True)
+
+    for it in items:
+        it["label"] = ("easy" if it["direct_ok"]
+                       else "hard" if it["plan_ok"] else "drop")
+        it["sol_code"] = it["direct_code"] if it["direct_ok"] else it["plan_code"]
+
+    for split in ("train", "test"):
+        sub = [it for it in items if it["split"] == split]
+        n = len(sub)
+        print(f"{split}: n={n}  direct={sum(i['direct_ok'] for i in sub)/n:.3f}  "
+              f"plan={sum(i['plan_ok'] for i in sub)/n:.3f}  "
+              f"easy={sum(i['label']=='easy' for i in sub)}  "
+              f"hard={sum(i['label']=='hard' for i in sub)}  "
+              f"drop={sum(i['label']=='drop' for i in sub)}", flush=True)
+
+    with open(OUT / "mbpp_data.json", "w") as f:
+        json.dump(items, f, indent=1)
+    print("wrote", OUT / "mbpp_data.json")
+
+
+if __name__ == "__main__":
+    main()