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>
60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
"""Render heatmaps (concept P by layer x position) and the layer profile."""
|
|
|
|
import os, sys
|
|
from pathlib import Path
|
|
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import torch
|
|
|
|
RES = Path(__file__).resolve().parent.parent / os.environ.get("JLENS_RESULTS", "results")
|
|
|
|
|
|
def plot_heat(path):
|
|
d = torch.load(path)
|
|
m, words, toks = d["map"].detach(), d["words"], d["tokens"]
|
|
fig, ax = plt.subplots(figsize=(min(16, 0.28 * m.shape[1] + 2), 6))
|
|
im = ax.imshow(m.numpy(), aspect="auto", origin="lower", cmap="magma",
|
|
vmin=0)
|
|
ax.set_ylabel("layer")
|
|
ax.set_xlabel("position")
|
|
if toks and len(toks) <= 60:
|
|
ax.set_xticks(range(len(toks)))
|
|
ax.set_xticklabels([t.replace("\n", "\\n") for t in toks],
|
|
rotation=90, fontsize=6)
|
|
ax.set_title(f"J-lens P({'/'.join(words)})")
|
|
fig.colorbar(im)
|
|
out = path.with_suffix(".png")
|
|
fig.tight_layout()
|
|
fig.savefig(out, dpi=130)
|
|
plt.close(fig)
|
|
print("wrote", out)
|
|
|
|
|
|
def plot_profile(path):
|
|
d = torch.load(path)
|
|
fig, ax1 = plt.subplots(figsize=(8, 4.5))
|
|
L = len(d["entropy"])
|
|
ax1.plot(range(L), d["entropy"], "k-", label="lens entropy")
|
|
ax1.set_xlabel("layer")
|
|
ax1.set_ylabel("entropy (nats)")
|
|
ax2 = ax1.twinx()
|
|
ax2.plot(range(L), d["agree_in"], "b--", label="top tok == current input")
|
|
ax2.plot(range(L), d["agree_next"], "r--", label="top tok == next output")
|
|
ax2.set_ylabel("agreement")
|
|
fig.legend(loc="upper center", ncol=3, fontsize=8)
|
|
fig.tight_layout()
|
|
out = path.with_suffix(".png")
|
|
fig.savefig(out, dpi=130)
|
|
plt.close(fig)
|
|
print("wrote", out)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
for p in sorted(RES.glob("heat_*.pt")):
|
|
plot_heat(p)
|
|
lp = RES / "layer_profile.pt"
|
|
if lp.exists():
|
|
plot_profile(lp)
|