| """Boot-prompt readout through the merged Huginn r=16 lens: per virtual layer, |
| transport h_v to the target and rank the prompt-absent answer token (Italy).""" |
|
|
| import torch |
| from jlens.lens import JacobianLens |
| from jlens.hooks import ActivationRecorder |
| from huginn_lens import HuginnLensModel, load_huginn |
|
|
| hf, tok = load_huginn() |
| hf = hf.cuda() |
| model = HuginnLensModel(hf, tok, num_steps=16) |
| lens = JacobianLens.load("jacobian_lens_huginn-0125_r16.pt") |
| print(model, flush=True) |
|
|
| prompt = "Fact: The currency used in the country shaped like a boot is" |
| ids = model.encode(prompt) |
| pos = ids.shape[1] - 1 |
| italy = tok.encode(" Italy")[-1] |
| print("italy token:", italy, repr(tok.decode([italy])), flush=True) |
|
|
| layers = list(range(model.n_layers - 1)) |
| with ActivationRecorder(model.layers, at=layers) as rec, torch.no_grad(): |
| out = model.forward(ids) |
|
|
| print(f"{'v':>3} {'kind':<10} {'top1':<16} {'italy_rank':>10}", flush=True) |
| for v in layers: |
| h = rec.activations[v][0, pos].float().cuda() |
| J = lens.jacobians[v].float().cuda() |
| logits = model.unembed((h @ J.T).to(torch.bfloat16)) |
| rank = int((logits > logits[italy]).sum().item()) + 1 |
| top1 = tok.decode([logits.argmax().item()]) |
| if v < model.n_prelude: |
| kind = "prelude" |
| elif v >= model.n_prelude + model.num_steps * model.n_core: |
| kind = "coda" |
| else: |
| kind = f"step{(v - model.n_prelude) // model.n_core}.L{(v - model.n_prelude) % model.n_core}" |
| print(f"{v:>3} {kind:<10} {top1!r:<16} {rank:>10}", flush=True) |
| print("PROBE_DONE", flush=True) |
|
|
