"""
visualize.py -- live visualizer for Wheeler-LM: watch the "wavefunction of the universe"
mixer as the model generates, token by token.
There is no colony and no space here -- the Wheeler-DeWitt block is TIMELESS, so the
right things to watch are the quantities the equation is actually about:
Captures per generated token (via instrument.Recorder, which WheelerDeWittBlock writes to):
* the TIMELIKE mode Psi_0 of each layer -- the emergent "clock" (volume/scale direction)
* the Hamiltonian-constraint residual |H| per layer -- how close to the physical H=0 surface
* the full K minisuperspace modes Psi for one layer -- the state on superspace
* per-layer block norm + attention + top predictions + token stream
Renders one self-contained HTML dashboard (pure inline SVG, no external scripts / CDN):
* "Emergent time" -- Psi_0 (the clock) traced across every generated token
* "Hamiltonian constraint |H|" -- traced across tokens (lower = more physical)
* "Minisuperspace state" -- the K modes at the current token, mode 0 (timelike) highlighted
* "|H| per layer" heat, attention, token stream, top predictions, trait activity
Usage:
python visualize.py --prompt "the universe" --tokens 60
"""
import argparse, json, os, sys, webbrowser
import torch
from model import QuazimotoLM, QuazimotoConfig
import instrument
PKG_DIR = os.path.dirname(os.path.abspath(__file__))
def find_ckpt(path):
if path and os.path.isfile(path):
return path
folder = (os.path.dirname(path) if path else os.path.join(PKG_DIR, "chkpt")) or "."
if not os.path.isdir(folder):
return None
pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")]
return max(pts, key=os.path.getmtime) if pts else None
def load_tokenizer(tok_dir):
sys.path.insert(0, tok_dir)
from spike_tokenizer import SpikeTokenizer
return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
@torch.no_grad()
def run_capture(model, cfg, tok, ids, n_tokens, temperature, top_k, device):
pl = cfg.n_layer // 2
rec = instrument.Recorder(phase_layer=pl, attn_layer=pl)
instrument.set_rec(rec)
idx = torch.tensor([ids], device=device)
try:
for _ in range(n_tokens):
cond = idx[:, -cfg.block_size:]
rec.begin()
logits, _, _ = model(cond)
lg = torch.nan_to_num(logits[:, -1, :].float(), nan=0.0,
posinf=1e4, neginf=-1e4) / max(temperature, 1e-6)
if top_k:
v = torch.topk(lg, min(top_k, lg.size(-1))).values
lg = lg.masked_fill(lg < v[:, [-1]], float("-inf"))
probs = torch.softmax(lg, dim=-1)
nxt = (torch.argmax(lg, -1, keepdim=True) if temperature <= 1e-3
else torch.multinomial(probs, 1))
top = torch.topk(torch.softmax(logits[:, -1].float(), -1), 5)
rec.end(token=int(nxt),
char=tok.decode([int(nxt)], skip_special_tokens=False),
top=[[int(t), round(float(p), 3)] for t, p in zip(top.indices[0], top.values[0])])
idx = torch.cat([idx, nxt], dim=1)
finally:
instrument.set_rec(None)
return rec.frames, idx[0].tolist()
def main():
p = argparse.ArgumentParser(description="Wheeler-LM emergent-time / constraint visualizer")
p.add_argument("--ckpt", default="")
p.add_argument("--tok_dir", default=PKG_DIR)
p.add_argument("--prompt", default="the universe")
p.add_argument("--tokens", type=int, default=60)
p.add_argument("--temperature", type=float, default=0.0)
p.add_argument("--top_k", type=int, default=40)
p.add_argument("--device", default="cpu")
p.add_argument("--out", default=os.path.join(PKG_DIR, "viz.html"))
p.add_argument("--no_open", action="store_true")
args = p.parse_args()
path = find_ckpt(args.ckpt)
if path is None:
print("No checkpoint found; pass --ckpt or train one."); return
ckpt = torch.load(path, map_location=args.device, weights_only=False)
cfg = QuazimotoConfig(**ckpt["family_config"])
model = QuazimotoLM(cfg); model.load_state_dict(ckpt["model"], strict=False)
model.to(args.device).eval()
tok = load_tokenizer(args.tok_dir)
print(f"loaded {path} (step {ckpt.get('step')}) | capturing {args.tokens} tokens ...")
ids = tok.encode(args.prompt, add_special_tokens=False)
frames, _ = run_capture(model, cfg, tok, ids, args.tokens,
args.temperature, args.top_k or None, args.device)
data = {
"meta": {"ckpt": os.path.basename(path), "step": ckpt.get("step"),
"n_layer": cfg.n_layer, "wdw_modes": cfg.wdw_modes,
"wdw_steps": cfg.wdw_steps, "phase_layer": cfg.n_layer // 2,
"prompt": args.prompt},
"prompt_chars": [tok.decode([t], skip_special_tokens=False) for t in ids],
"frames": frames,
}
html = HTML_TEMPLATE.replace("/*__DATA__*/", json.dumps(data))
with open(args.out, "w", encoding="utf-8") as f:
f.write(html)
print(f"wrote {args.out} ({os.path.getsize(args.out)//1024} KB)")
if not args.no_open:
webbrowser.open("file://" + os.path.abspath(args.out))
HTML_TEMPLATE = r"""
Wheeler-LM live
Trait activity (this token)
Emergent time — timelike mode Ψ₀ (layer ) across tokens
Hamiltonian constraint |H| across tokens (lower = more physical, HΨ=0)
Minisuperspace state — K modes Ψ at this token (orange = timelike)
|H| per layer (this token)
Attention — layer (last token → context)
"""
if __name__ == "__main__":
main()