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,104 @@
+"""Sampled relabeling of test-split difficulty buckets (protocol item:
+outcome-selection fix).
+
+Greedy labeling couples bucket selection to the same coin flips as the k=0
+eval (circular: a bucket defined by baseline failure shows baseline ~0%).
+Fix: labels from 3 temperature-sampled direct attempts, independent of the
+greedy eval — hard_s = 0/3 sampled direct correct AND (plan/CoT reachable);
+easy_s = >=2/3 correct; else mid_s. Adds 'label_sampled' to the data JSONs.
+"""
+
+import json
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import torch
+
+from loop_common import DIRECT_SUFFIX as GSM_SUFFIX
+from loop_common import chat_prompt, last_number, num_eq
+from prep_mbpp import (DIRECT_SUFFIX as MBPP_SUFFIX, 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"
+N_SAMPLES = 3
+TEMP = 0.8
+
+
+@torch.no_grad()
+def sample_batch(model, tok, prompts, max_new, batch=24, seed=0):
+    outs = [[] for _ in prompts]
+    for s in range(N_SAMPLES):
+        torch.manual_seed(1000 + s + seed)
+        for i in range(0, len(prompts), batch):
+            chunk = prompts[i : i + batch]
+            enc = tok(chunk, return_tensors="pt", padding=True,
+                      add_special_tokens=False).to("cuda")
+            gen = model.generate(**enc, max_new_tokens=max_new, do_sample=True,
+                                 temperature=TEMP, top_p=0.95,
+                                 pad_token_id=tok.pad_token_id or 0)
+            for j in range(len(chunk)):
+                outs[i + j].append(
+                    tok.decode(gen[j, enc["input_ids"].shape[1]:],
+                               skip_special_tokens=True))
+        print(f"  sample {s+1}/{N_SAMPLES} done", flush=True)
+    return outs
+
+
+def relabel(items, n_correct, reachable_key):
+    for it, nc in zip(items, n_correct):
+        if nc == 0:
+            it["label_sampled"] = ("hard" if it[reachable_key] else "drop")
+        elif nc >= 2:
+            it["label_sampled"] = "easy"
+        else:
+            it["label_sampled"] = "mid"
+
+
+def main():
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+
+    # --- MBPP test ---
+    mbpp = json.load(open(OUT / "mbpp_data.json"))
+    mtest = [it for it in mbpp if it["split"] == "test"]
+    print(f"MBPP: sampling {len(mtest)} items x{N_SAMPLES}", flush=True)
+    gens = sample_batch(model, tok,
+                        [mbpp_prompt(tok, it, MBPP_SUFFIX) for it in mtest],
+                        max_new=220)
+    nc = []
+    with ThreadPoolExecutor(8) as ex:
+        for it, gs in zip(mtest, gens):
+            oks = list(ex.map(lambda g: run_tests(extract_code(g), it), gs))
+            nc.append(sum(oks))
+    relabel(mtest, nc, "plan_ok")
+    json.dump(mbpp, open(OUT / "mbpp_data.json", "w"), indent=1)
+    dist = {l: sum(it.get("label_sampled") == l for it in mtest)
+            for l in ("easy", "mid", "hard", "drop")}
+    agree = sum(it["label"] == it.get("label_sampled") for it in mtest
+                if it.get("label_sampled") != "mid")
+    print(f"MBPP sampled labels: {dist}  (greedy-agreement excl. mid: "
+          f"{agree}/{sum(1 for it in mtest if it.get('label_sampled') != 'mid')})",
+          flush=True)
+
+    # --- GSM8K test ---
+    gsm = json.load(open(OUT / "star_data.json"))
+    gtest = [it for it in gsm if it["split"] == "test"]
+    print(f"GSM: sampling {len(gtest)} items x{N_SAMPLES}", flush=True)
+    gens = sample_batch(model, tok,
+                        [chat_prompt(tok, it["question"], GSM_SUFFIX)
+                         for it in gtest], max_new=10)
+    nc = [sum(num_eq(last_number(g), it["gold"]) for g in gs)
+          for it, gs in zip(gtest, gens)]
+    relabel(gtest, nc, "cot_ok")
+    json.dump(gsm, open(OUT / "star_data.json", "w"), indent=1)
+    dist = {l: sum(it.get("label_sampled") == l for it in gtest)
+            for l in ("easy", "mid", "hard", "drop")}
+    print(f"GSM sampled labels: {dist}", flush=True)
+
+
+if __name__ == "__main__":
+    main()