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>
66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
"""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()
|