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>
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
"""J-lens sharpening battery on MBPP prompts: does the trained loop
|
|
concentrate the workspace readout while reading the problem?
|
|
|
|
For N test prompts, at each k: J-lens distribution at L30, last prompt
|
|
position -> entropy + top-1 probability. Trained vs untrained adapter.
|
|
"""
|
|
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from loop_common import BandLooper, MergeAdapter
|
|
from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import JLens, load_model # noqa: E402
|
|
|
|
OUT = Path(__file__).resolve().parent.parent / "results-loop"
|
|
ROOT = OUT.parent
|
|
N_ITEMS = 20
|
|
KS = (0, 1, 2, 4)
|
|
|
|
|
|
@torch.no_grad()
|
|
def battery(looper, adapter, lens, tok, items):
|
|
rows = []
|
|
for it in items:
|
|
ids = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX), return_tensors="pt",
|
|
add_special_tokens=False)["input_ids"].cuda()
|
|
calls, _ = looper.capture(ids, logits_to_keep=1)
|
|
e = looper._hin[looper.l0]
|
|
s = looper.band(e, calls)
|
|
state = {0: s}
|
|
for k in range(1, max(KS) + 1):
|
|
s = looper.band(adapter(e, s), calls)
|
|
state[k] = s
|
|
for k in KS:
|
|
probs = torch.softmax(lens._readout(
|
|
state[k][0, -1].float() @ lens.Jbar[looper.l1].T).float(), -1)
|
|
ent = -(probs * (probs + 1e-12).log()).sum().item()
|
|
rows.append({"task_id": it["task_id"], "k": k, "entropy": ent,
|
|
"top1": probs.max().item()})
|
|
return rows
|
|
|
|
|
|
def main():
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
looper = BandLooper(model)
|
|
jbar = torch.load(ROOT / "results" / "jbar.pt", map_location="cuda")
|
|
lens = JLens(model, tok, jbar["Jbar"].float())
|
|
items = [it for it in json.load(open(OUT / "mbpp_data.json"))
|
|
if it["split"] == "test"][:N_ITEMS]
|
|
|
|
res = {}
|
|
for tag, path in (("untrained", None),
|
|
("trained", OUT / "adapter_code_e399.pt")):
|
|
adapter = MergeAdapter().cuda()
|
|
if path:
|
|
adapter.load_state_dict(torch.load(path, map_location="cuda"))
|
|
rows = battery(looper, adapter, lens, tok, items)
|
|
res[tag] = rows
|
|
for k in KS:
|
|
sel = [r for r in rows if r["k"] == k]
|
|
ent = sum(r["entropy"] for r in sel) / len(sel)
|
|
top = sum(r["top1"] for r in sel) / len(sel)
|
|
print(f"[{tag}] k={k}: mean_entropy={ent:.3f} mean_top1={top:.3f}",
|
|
flush=True)
|
|
json.dump(res, open(OUT / "lens_battery.json", "w"), indent=1)
|
|
print("wrote", OUT / "lens_battery.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|