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,65 @@
+"""Re-run the MBPP plan pass on direct-fail items with a non-truncating budget.
+
+The first plan pass (max_new=380) truncated ~all outputs before the code:
+E2B writes verbose plans. Fix: terse-plan prompt + 700-token budget. Only
+items with label=='drop' need re-labeling (easy is decided by the direct pass).
+"""
+
+import json
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import torch
+
+from prep_mbpp import batch_generate, extract_code, mbpp_prompt, run_tests
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(__file__).resolve().parent.parent / "results-loop"
+
+TERSE_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.")
+
+
+def main():
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+    data = json.load(open(OUT / "mbpp_data.json"))
+    redo = [it for it in data if it["label"] == "drop"]
+    print(f"re-running plan pass on {len(redo)} direct-fail items", flush=True)
+
+    t0 = time.time()
+    gens = batch_generate(model, tok,
+                          [mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX)
+                           for it in redo], max_new_tokens=700, batch_size=16)
+    codes = [extract_code(g) for g in gens]
+    n_trunc = sum(1 for g in gens if len(tok(g)["input_ids"]) >= 695)
+    with ThreadPoolExecutor(8) as ex:
+        oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
+                          zip(codes, redo)))
+    for it, c, ok in zip(redo, codes, oks):
+        it["plan_ok"] = bool(ok)
+        it["plan_code"] = c if ok else None
+        it["label"] = "hard" if ok else "drop"
+        it["sol_code"] = it["direct_code"] if it["direct_ok"] else it["plan_code"]
+    print(f"done in {time.time()-t0:.0f}s  plan-pass on fails: "
+          f"{sum(oks)}/{len(redo)}  still-truncated={n_trunc}", flush=True)
+
+    for split in ("train", "test"):
+        sub = [it for it in data if it["split"] == split]
+        n = len(sub)
+        print(f"{split}: n={n}  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(data, f, indent=1)
+    print("wrote", OUT / "mbpp_data.json")
+
+
+if __name__ == "__main__":
+    main()