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>
133 lines
5.0 KiB
Python
133 lines
5.0 KiB
Python
"""STaR-style difficulty labeling of MBPP with the frozen base model.
|
|
|
|
Two passes per item: direct code generation vs plan-first-then-code, each
|
|
executed against MBPP's unit tests (sandboxed subprocess). Labels:
|
|
easy (direct passes), hard (plan-only passes), drop (neither). The model's
|
|
own passing code is kept as the training target (in-distribution supervision).
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from datasets import load_dataset
|
|
|
|
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"))
|
|
OUT.mkdir(exist_ok=True)
|
|
|
|
DIRECT_SUFFIX = ("\n\nWrite only the Python function in a ```python code "
|
|
"block. No explanation.")
|
|
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.")
|
|
|
|
CODE_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.S)
|
|
|
|
|
|
def mbpp_prompt(tok, item, suffix):
|
|
tests = "\n".join(item["test_list"])
|
|
msg = (f"{item['text']}\nYour code should pass these tests:\n\n{tests}"
|
|
f"{suffix}")
|
|
return tok.apply_chat_template([{"role": "user", "content": msg}],
|
|
tokenize=False, add_generation_prompt=True)
|
|
|
|
|
|
def extract_code(text):
|
|
m = CODE_RE.findall(text)
|
|
return m[-1].strip() if m else None
|
|
|
|
|
|
def run_tests(code, item, timeout=10):
|
|
if not code:
|
|
return False
|
|
script = (item.get("test_setup_code") or "") + "\n" + code + "\n" + \
|
|
"\n".join(item["test_list"])
|
|
try:
|
|
with tempfile.TemporaryDirectory() as td:
|
|
r = subprocess.run([sys.executable, "-c", script], cwd=td,
|
|
capture_output=True, timeout=timeout)
|
|
return r.returncode == 0
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
return False
|
|
|
|
|
|
@torch.no_grad()
|
|
def batch_generate(model, tok, prompts, max_new_tokens, batch_size=24):
|
|
outs = []
|
|
for i in range(0, len(prompts), batch_size):
|
|
chunk = prompts[i : i + batch_size]
|
|
enc = tok(chunk, return_tensors="pt", padding=True,
|
|
add_special_tokens=False).to("cuda")
|
|
gen = model.generate(**enc, max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
pad_token_id=tok.pad_token_id or 0)
|
|
for j in range(len(chunk)):
|
|
outs.append(tok.decode(gen[j, enc["input_ids"].shape[1]:],
|
|
skip_special_tokens=True))
|
|
print(f" {min(i+batch_size, len(prompts))}/{len(prompts)}", flush=True)
|
|
return outs
|
|
|
|
|
|
def main():
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
tok.padding_side = "left"
|
|
ds = load_dataset("google-research-datasets/mbpp", "full")
|
|
|
|
items = []
|
|
for split, tag in (("train", "train"), ("validation", "train"),
|
|
("test", "test")):
|
|
for row in ds[split]:
|
|
items.append({"split": tag, "task_id": row["task_id"],
|
|
"text": row["text"], "test_list": row["test_list"],
|
|
"test_setup_code": row["test_setup_code"]})
|
|
print(f"{len(items)} items", flush=True)
|
|
|
|
for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 220),
|
|
("plan", PLAN_SUFFIX, 700)):
|
|
t0 = time.time()
|
|
print(f"{tag} pass", flush=True)
|
|
gens = batch_generate(model, tok,
|
|
[mbpp_prompt(tok, it, suffix) for it in items], mx)
|
|
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)))
|
|
for it, c, ok in zip(items, codes, oks):
|
|
it[f"{tag}_ok"] = bool(ok)
|
|
it[f"{tag}_code"] = c if ok else None
|
|
print(f"{tag} pass done in {time.time()-t0:.0f}s "
|
|
f"pass@1={sum(oks)/len(items):.3f}", flush=True)
|
|
|
|
for it in items:
|
|
it["label"] = ("easy" if it["direct_ok"]
|
|
else "hard" if it["plan_ok"] else "drop")
|
|
it["sol_code"] = it["direct_code"] if it["direct_ok"] else it["plan_code"]
|
|
|
|
for split in ("train", "test"):
|
|
sub = [it for it in items if it["split"] == split]
|
|
n = len(sub)
|
|
print(f"{split}: n={n} direct={sum(i['direct_ok'] for i in sub)/n:.3f} "
|
|
f"plan={sum(i['plan_ok'] for i in sub)/n:.3f} "
|
|
f"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(items, f, indent=1)
|
|
print("wrote", OUT / "mbpp_data.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|