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>
104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
"""STaR-style difficulty labeling of GSM8K with the frozen base model.
|
|
|
|
Two passes per item with frozen gemma-4-E2B-it:
|
|
direct: answer with no CoT -> correct = "easy"
|
|
cot: step-by-step -> correct (and direct wrong) = "hard"
|
|
Items the model cannot solve even with CoT are dropped from training
|
|
(unreachable supervision) but kept in the test split for eval.
|
|
|
|
Output: results-loop/star_data.json
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from datasets import load_dataset
|
|
|
|
from loop_common import (COT_SUFFIX, DIRECT_SUFFIX, chat_prompt, gold_answer,
|
|
last_number, num_eq)
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import load_model # noqa: E402
|
|
|
|
N_TRAIN, N_TEST = 1024, 256
|
|
OUT = Path(os.environ.get("LOOP_OUT",
|
|
Path(__file__).resolve().parent.parent / "results-loop"))
|
|
OUT.mkdir(exist_ok=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def batch_generate(model, tok, prompts, max_new_tokens, batch_size):
|
|
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("openai/gsm8k", "main")
|
|
|
|
rng = random.Random(0)
|
|
tr_idx = rng.sample(range(len(ds["train"])), N_TRAIN)
|
|
te_idx = rng.sample(range(len(ds["test"])), N_TEST)
|
|
|
|
items = []
|
|
for split, idxs in (("train", tr_idx), ("test", te_idx)):
|
|
for i in idxs:
|
|
row = ds[split][i]
|
|
items.append({"split": split, "idx": i, "question": row["question"],
|
|
"gold": gold_answer(row["answer"])})
|
|
|
|
t0 = time.time()
|
|
print(f"direct pass ({len(items)} items)", flush=True)
|
|
direct = batch_generate(
|
|
model, tok, [chat_prompt(tok, it["question"], DIRECT_SUFFIX) for it in items],
|
|
max_new_tokens=10, batch_size=64)
|
|
print(f"direct pass done in {time.time()-t0:.0f}s", flush=True)
|
|
|
|
t0 = time.time()
|
|
print("cot pass", flush=True)
|
|
cot = batch_generate(
|
|
model, tok, [chat_prompt(tok, it["question"], COT_SUFFIX) for it in items],
|
|
max_new_tokens=320, batch_size=32)
|
|
print(f"cot pass done in {time.time()-t0:.0f}s", flush=True)
|
|
|
|
for it, d, c in zip(items, direct, cot):
|
|
it["direct_pred"] = last_number(d)
|
|
it["direct_ok"] = num_eq(it["direct_pred"], it["gold"])
|
|
it["cot_pred"] = last_number(c)
|
|
it["cot_ok"] = num_eq(it["cot_pred"], it["gold"])
|
|
it["label"] = ("easy" if it["direct_ok"]
|
|
else "hard" if it["cot_ok"] else "drop")
|
|
|
|
for split in ("train", "test"):
|
|
sub = [it for it in items if it["split"] == split]
|
|
n = len(sub)
|
|
print(f"{split}: n={n} direct_acc={sum(i['direct_ok'] for i in sub)/n:.3f} "
|
|
f"cot_acc={sum(i['cot_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 / "star_data.json", "w") as f:
|
|
json.dump(items, f, indent=1)
|
|
print("wrote", OUT / "star_data.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|