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>
55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
"""Explicit-planning reference: frozen model, plan-first prompt, MBPP test.
|
|
|
|
The number latent looping is compared against (equal-or-more FLOPs: ~200-400
|
|
visible plan tokens vs k band passes over the prompt)."""
|
|
|
|
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
|
|
from prep_mbpp_fix import TERSE_PLAN_SUFFIX
|
|
|
|
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"
|
|
|
|
|
|
def main():
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
tok.padding_side = "left"
|
|
items = [it for it in json.load(open(OUT / "mbpp_data.json"))
|
|
if it["split"] == "test"]
|
|
t0 = time.time()
|
|
gens = batch_generate(model, tok,
|
|
[mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX)
|
|
for it in items], max_new_tokens=700, batch_size=16)
|
|
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)))
|
|
by = {}
|
|
for it, ok in zip(items, oks):
|
|
d = by.setdefault(it["label"], [0, 0])
|
|
d[0] += ok
|
|
d[1] += 1
|
|
per_item = [{"task_id": it["task_id"], "ok": bool(ok)}
|
|
for it, ok in zip(items, oks)]
|
|
res = {"acc": sum(oks) / len(items),
|
|
"by_label": {l: c / n for l, (c, n) in by.items()},
|
|
"per_item": per_item}
|
|
print(f"plan-first pass@1={res['acc']:.3f} "
|
|
f"by_label={ {l: round(v,3) for l,v in res['by_label'].items()} }"
|
|
f" ({time.time()-t0:.0f}s)", flush=True)
|
|
json.dump(res, open(OUT / "eval_plan_baseline.json", "w"), indent=1)
|
|
print("wrote", OUT / "eval_plan_baseline.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|