NilsandClaude Fable 5
ef9c08966c
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>
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
"""Blocksworld STaR labeling with the frozen model (direct vs CoT plan)."""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from bw_common import make_dataset, verify_plan
|
|
from prep_mbpp import batch_generate
|
|
|
|
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"))
|
|
|
|
DIRECT_SUFFIX = "\n\nGive only the numbered list of moves, nothing else."
|
|
COT_SUFFIX = ("\n\nFirst think step by step about which blocks must move "
|
|
"and in what order (briefly), then give the numbered list of "
|
|
"moves.")
|
|
|
|
|
|
def chat(tok, it, suffix):
|
|
return tok.apply_chat_template(
|
|
[{"role": "user", "content": it["question"] + suffix}],
|
|
tokenize=False, add_generation_prompt=True)
|
|
|
|
|
|
def main():
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
tok.padding_side = "left"
|
|
items = make_dataset(n_train=400, n_test=200, seed=0)
|
|
print(f"{len(items)} instances", flush=True)
|
|
|
|
for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 200),
|
|
("cot", COT_SUFFIX, 500)):
|
|
t0 = time.time()
|
|
gens = batch_generate(model, tok,
|
|
[chat(tok, it, suffix) for it in items],
|
|
max_new_tokens=mx, batch_size=24)
|
|
oks = [verify_plan(it, g) for it, g in zip(items, gens)]
|
|
for it, g, ok in zip(items, gens, oks):
|
|
it[f"{tag}_ok"] = bool(ok)
|
|
it[f"{tag}_plan"] = g if ok else None
|
|
print(f"{tag}: acc={sum(oks)/len(items):.3f} "
|
|
f"({time.time()-t0:.0f}s)", flush=True)
|
|
|
|
for it in items:
|
|
it["label"] = ("easy" if it["direct_ok"]
|
|
else "hard" if it["cot_ok"] else "drop")
|
|
it["sol_plan"] = (it["direct_plan"] if it["direct_ok"]
|
|
else it["cot_plan"])
|
|
for split in ("train", "test"):
|
|
sub = [it for it in items if it["split"] == split]
|
|
print(f"{split}: 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)
|
|
json.dump(items, open(OUT / "bw_data.json", "w"), indent=1)
|
|
print("wrote", OUT / "bw_data.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|