| """Deliberation-trace stage 1: per-recurrence J-lens readout over a riddle set |
| with known plausible-but-wrong candidates (full 16-lens tiling). |
| |
| Per riddle, one recorded forward. At every virtual layer: rank of the |
| expected answer and each trap candidate, projection of the state onto each |
| token's lens direction (plus fixed control tokens), and the top-8 decoded |
| tokens. Plus a 12-token greedy continuation for the model's actual answer. |
| """ |
|
|
| import json |
|
|
| import torch |
|
|
| import ws_common as W |
|
|
| torch.set_grad_enabled(False) |
|
|
| ctx = W.load_ctx("huginn", full_depth=True, huginn_full_tiling=True) |
| V = 67 |
|
|
| RIDDLES = [ |
| ("boot_currency", "Fact: The currency used in the country shaped like a boot is the", |
| "euro", ["lira", "Italy", "Italian"]), |
| ("eiffel_pre2002", "The currency used in the country of the Eiffel Tower before 2002 was the", |
| "franc", ["euro", "France"]), |
| ("sushi_currency", "The currency used in the country famous for sushi is the", |
| "yen", ["yuan", "won"]), |
| ("kangaroo_capital", "The capital of the country where kangaroos live is", |
| "Canberra", ["Sydney", "Melbourne"]), |
| ("maple_capital", "The capital of the country famous for maple syrup is", |
| "Ottawa", ["Toronto", "Montreal"]), |
| ("liberty_capital", "The capital of the country home to the Statue of Liberty is", |
| "Washington", ["New", "York", "Paris"]), |
| ("largest_planet", "The largest planet in the solar system is", |
| "Jupiter", ["Saturn", "Neptune"]), |
| ("closest_planet", "The planet closest to the sun is", |
| "Mercury", ["Venus", "Mars"]), |
| ("fastest_animal", "The fastest land animal is the", |
| "cheetah", ["lion", "leopard"]), |
| ("jungle_king", "The animal known as the king of the jungle is the", |
| "lion", ["tiger", "elephant"]), |
| ("tallest_mountain", "The tallest mountain in the world is Mount", |
| "Everest", ["Fuji", "Kilimanjaro"]), |
| ("relativity", "The scientist who developed the theory of relativity was", |
| "Einstein", ["Newton", "Darwin"]), |
| ] |
|
|
| CONTROLS = ["table", "music", "seven", "window", "doctor", "river", |
| "yellow", "engine"] |
|
|
|
|
| def greedy(prompt, n=12): |
| ids = W.encode(ctx, prompt) |
| outs = [] |
| for _ in range(n): |
| logits, _ = W.forward(ctx, ids, layers=[]) |
| t = int(logits[-1].argmax()) |
| outs.append(t) |
| ids = torch.cat([ids, torch.tensor([[t]], device=ids.device)], dim=1) |
| return ctx.tok.decode(outs) |
|
|
|
|
| out = {"riddles": []} |
| ctl_toks = [W.first_token(ctx, w) for w in CONTROLS] |
| for name, prompt, answer, traps in RIDDLES: |
| toks = {answer: W.first_token(ctx, answer)} |
| for t in traps: |
| toks[t] = W.first_token(ctx, t) |
| logits, acts = W.forward(ctx, W.encode(ctx, prompt), layers=list(range(V))) |
| rec = {"name": name, "prompt": prompt, "answer": answer, "traps": traps, |
| "argmax": ctx.tok.decode([int(logits[-1].argmax())]), |
| "gen": greedy(prompt), |
| "ranks": {w: [] for w in toks}, "proj": {w: [] for w in toks}, |
| "ctl_proj": [], "top8": []} |
| for v in range(V): |
| h = acts[v][-1] |
| lg = ctx.model.unembed(h @ ctx.Jt[v]).float() |
| for w, t in toks.items(): |
| rec["ranks"][w].append(int((lg > lg[t]).sum().item()) + 1) |
| d = W.direction(ctx, v, t) |
| rec["proj"][w].append(float(h @ d)) |
| rec["ctl_proj"].append([float(h @ W.direction(ctx, v, t)) |
| for t in ctl_toks]) |
| rec["top8"].append([ctx.tok.decode([t]) for t in |
| lg.topk(8).indices.tolist()]) |
| out["riddles"].append(rec) |
| print(f"done {name}: argmax={rec['argmax']!r} gen={rec['gen'][:40]!r}", |
| flush=True) |
|
|
| json.dump(out, open("delib_trace.json", "w")) |
| print("wrote delib_trace.json", flush=True) |
| print("DONE_DELIB", flush=True) |
|
|
