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>
68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
"""Estimate J_l = E_{prompt, t, t'>=t}[dh_final,t'/dh_l,t] over a pretraining-like corpus."""
|
|
|
|
import argparse, json, sys, time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import load_model, prompt_jacobian_pairsum
|
|
|
|
|
|
def corpus_texts(n, min_chars=400):
|
|
from datasets import load_dataset
|
|
ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT",
|
|
split="train", streaming=True)
|
|
got = 0
|
|
for ex in ds:
|
|
t = ex["text"].strip()
|
|
if len(t) >= min_chars:
|
|
yield t
|
|
got += 1
|
|
if got >= n:
|
|
return
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--n-prompts", type=int, default=256)
|
|
ap.add_argument("--seq-len", type=int, default=64)
|
|
ap.add_argument("--chunk", type=int, default=128)
|
|
ap.add_argument("--out", default="results/jbar.pt")
|
|
ap.add_argument("--dtype", default="float32")
|
|
args = ap.parse_args()
|
|
|
|
model, tok = load_model(dtype=getattr(torch, args.dtype))
|
|
L = len(model.model.language_model.layers)
|
|
d = model.config.get_text_config().hidden_size
|
|
|
|
Jsum = torch.zeros(L, d, d, device="cuda", dtype=torch.float32)
|
|
pairs_total = 0
|
|
t0 = time.time()
|
|
for i, text in enumerate(corpus_texts(args.n_prompts)):
|
|
ids = tok(text, return_tensors="pt", truncation=True,
|
|
max_length=args.seq_len)["input_ids"].cuda()
|
|
if ids.shape[1] < args.seq_len:
|
|
continue
|
|
J, n_pairs = prompt_jacobian_pairsum(model, ids, chunk=args.chunk)
|
|
Jsum += J
|
|
pairs_total += n_pairs
|
|
if i % 5 == 0 or i == args.n_prompts - 1:
|
|
el = time.time() - t0
|
|
print(f"[{i+1}/{args.n_prompts}] {el:.0f}s ({el/(i+1):.1f}s/prompt)",
|
|
flush=True)
|
|
if i % 50 == 49: # checkpoint
|
|
torch.save({"Jbar": (Jsum / pairs_total).cpu(),
|
|
"n_prompts": i + 1, "seq_len": args.seq_len},
|
|
args.out + ".ckpt")
|
|
|
|
Jbar = (Jsum / pairs_total).cpu()
|
|
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save({"Jbar": Jbar, "n_prompts": args.n_prompts,
|
|
"seq_len": args.seq_len, "pairs": pairs_total}, args.out)
|
|
print("saved", args.out)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|