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>
127 lines
4.9 KiB
Python
127 lines
4.9 KiB
Python
"""Per-prompt gate: probe on the k=0 workspace state predicts 'will looping help?'
|
|
|
|
1. Features: L30 residual at the last prompt position (plain forward), MBPP
|
|
train items; labels easy(0)/hard(1) from the STaR pass (free supervision).
|
|
2. Logistic probe (d=1536 -> 1), class-balanced.
|
|
3. Gated eval on MBPP test: predicted-easy -> k=0, predicted-hard -> k=4 with
|
|
the dedicated loop adapter. Reports gated pass@1 vs uniform-k references.
|
|
"""
|
|
|
|
import json
|
|
import sys
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from loop_common import BandLooper, MergeAdapter
|
|
from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import ResidualCapture, load_model # noqa: E402
|
|
|
|
OUT = Path(__file__).resolve().parent.parent / "results-loop"
|
|
LAYER = 30
|
|
K_HARD = 4
|
|
|
|
|
|
@torch.no_grad()
|
|
def collect_states(model, tok, items, batch=16):
|
|
feats = []
|
|
for i in range(0, len(items), batch):
|
|
chunk = items[i : i + batch]
|
|
enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk],
|
|
return_tensors="pt", padding=True,
|
|
add_special_tokens=False).to("cuda")
|
|
with ResidualCapture(model, layers=[LAYER]) as cap:
|
|
model(**enc, use_cache=False, logits_to_keep=1)
|
|
h = cap.acts[LAYER] # (B, T, d); left padding -> last pos is prompt end
|
|
feats.append(h[:, -1].float().cpu())
|
|
return torch.cat(feats)
|
|
|
|
|
|
def fit_probe(X, y, epochs=300, lr=0.05):
|
|
mu, sd = X.mean(0), X.std(0) + 1e-6
|
|
Xn = (X - mu) / sd
|
|
w = torch.zeros(X.shape[1], requires_grad=True)
|
|
b = torch.zeros(1, requires_grad=True)
|
|
opt = torch.optim.Adam([w, b], lr=lr)
|
|
pos_w = (y == 0).sum() / max(1, (y == 1).sum())
|
|
for _ in range(epochs):
|
|
z = Xn @ w + b
|
|
loss = torch.nn.functional.binary_cross_entropy_with_logits(
|
|
z, y.float(), pos_weight=pos_w)
|
|
opt.zero_grad()
|
|
loss.backward()
|
|
opt.step()
|
|
return w.detach(), b.detach(), mu, sd
|
|
|
|
|
|
def main():
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
tok.padding_side = "left"
|
|
looper = BandLooper(model)
|
|
adapter = MergeAdapter().cuda()
|
|
adapter.load_state_dict(torch.load(OUT / "adapter_code_e399.pt",
|
|
map_location="cuda"))
|
|
|
|
data = json.load(open(OUT / "mbpp_data.json"))
|
|
train = [it for it in data if it["split"] == "train"
|
|
and it["label"] in ("easy", "hard")]
|
|
test = [it for it in data if it["split"] == "test"][:250]
|
|
|
|
print(f"collecting states: {len(train)} train, {len(test)} test", flush=True)
|
|
Xtr = collect_states(model, tok, train)
|
|
ytr = torch.tensor([it["label"] == "hard" for it in train]).long()
|
|
Xte = collect_states(model, tok, test)
|
|
|
|
w, b, mu, sd = fit_probe(Xtr, ytr)
|
|
ptr = torch.sigmoid(((Xtr - mu) / sd) @ w + b)
|
|
acc_tr = ((ptr > 0.5).long() == ytr).float().mean()
|
|
pte = torch.sigmoid(((Xte - mu) / sd) @ w + b)
|
|
pred_hard = (pte > 0.5).tolist()
|
|
yte = [it["label"] == "hard" for it in test]
|
|
tp = sum(p and t for p, t in zip(pred_hard, yte))
|
|
print(f"probe: train_acc={acc_tr:.3f} test: pred_hard="
|
|
f"{sum(pred_hard)} (true hard={sum(yte)}, tp={tp})", flush=True)
|
|
|
|
# gated generation: k=0 for predicted-easy, K_HARD for predicted-hard
|
|
codes = [None] * len(test)
|
|
for k, sel in ((0, [i for i, ph in enumerate(pred_hard) if not ph]),
|
|
(K_HARD, [i for i, ph in enumerate(pred_hard) if ph])):
|
|
for i in range(0, len(sel), 8):
|
|
idxs = sel[i : i + 8]
|
|
chunk = [test[j] for j in idxs]
|
|
enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk],
|
|
return_tensors="pt", padding=True,
|
|
add_special_tokens=False).to("cuda")
|
|
gen = looper.generate_frozen_prompt(
|
|
adapter, tok, enc["input_ids"], k, max_new_tokens=220,
|
|
attention_mask=enc["attention_mask"])
|
|
for jj, j in enumerate(idxs):
|
|
codes[j] = extract_code(
|
|
tok.decode(gen[jj, enc["input_ids"].shape[1]:],
|
|
skip_special_tokens=True))
|
|
with ThreadPoolExecutor(8) as ex:
|
|
oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
|
|
zip(codes, test)))
|
|
acc = sum(oks) / len(test)
|
|
by = {}
|
|
for it, ok in zip(test, oks):
|
|
d = by.setdefault(it["label"], [0, 0])
|
|
d[0] += ok
|
|
d[1] += 1
|
|
print(f"GATED pass@1={acc:.3f} "
|
|
f"by_label={ {l: round(c/n,3) for l,(c,n) in by.items()} }", flush=True)
|
|
json.dump({"acc": acc,
|
|
"by_label": {l: c / n for l, (c, n) in by.items()},
|
|
"probe_train_acc": float(acc_tr),
|
|
"pred_hard_test": int(sum(pred_hard))},
|
|
open(OUT / "eval_gated.json", "w"), indent=1)
|
|
print("wrote", OUT / "eval_gated.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|