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>
105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
"""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()
|