Wheeler-Chat / visualize.py
Quazim0t0's picture
Upload folder using huggingface_hub
f2da4d6 verified
Raw
History Blame Contribute Delete
14.4 kB
"""
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"""<!doctype html>
<html><head><meta charset="utf-8"><title>Wheeler-LM live</title>
<style>
:root{--bg:#0d1117;--panel:#161b22;--ink:#e6edf3;--mut:#8b949e;--ac:#58a6ff;--hot:#f0883e;--spore:#3fb950;--red:#f85149}
*{box-sizing:border-box} body{margin:0;background:var(--bg);color:var(--ink);font:13px/1.4 ui-monospace,Menlo,Consolas,monospace}
header{padding:10px 16px;background:var(--panel);border-bottom:1px solid #30363d;display:flex;gap:16px;align-items:center;flex-wrap:wrap}
h1{font-size:15px;margin:0;color:var(--ac)} .mut{color:var(--mut)}
#wrap{display:grid;grid-template-columns:320px 1fr 300px;gap:12px;padding:12px}
.panel{background:var(--panel);border:1px solid #30363d;border-radius:8px;padding:10px}
.panel h2{font-size:12px;margin:0 0 8px;color:var(--mut);text-transform:uppercase;letter-spacing:.5px}
#stream{line-height:1.9;max-height:120px;overflow:auto}
.tk{padding:1px 2px;border-radius:3px;cursor:pointer;white-space:pre}
.tk.prompt{color:var(--mut)} .tk.cur{background:var(--ac);color:#000} .tk:hover{background:#30363d}
.ctrl{display:flex;gap:10px;align-items:center} input[type=range]{width:260px}
button{background:#21262d;color:var(--ink);border:1px solid #30363d;border-radius:6px;padding:4px 10px;cursor:pointer}
button:hover{border-color:var(--ac)} svg{display:block;width:100%;height:auto}
.bar{height:10px;background:#21262d;border-radius:5px;overflow:hidden;margin:2px 0}.bar>div{height:100%}
table{border-collapse:collapse;width:100%}td{padding:1px 4px}.lbl{color:var(--mut);width:52px}
.grid{display:grid;gap:3px}.cell{height:15px;border-radius:2px}.small{font-size:11px;color:var(--mut)}
</style></head><body>
<header><h1>Wheeler-LM &middot; wavefunction of the universe</h1><span class="mut" id="meta"></span>
<span style="flex:1"></span>
<div class="ctrl"><button id="play">▶ play</button><input type="range" id="seek" min="0" value="0"><span id="pos" class="mut"></span></div>
</header>
<div id="wrap">
<div style="display:flex;flex-direction:column;gap:12px">
<div class="panel"><h2>Token stream</h2><div id="stream"></div></div>
<div class="panel"><h2>Trait activity (this token)</h2><div id="traits"></div></div>
<div class="panel"><h2>Top predictions</h2><div id="top"></div></div>
</div>
<div style="display:flex;flex-direction:column;gap:12px">
<div class="panel"><h2>Emergent time &mdash; timelike mode &Psi;&#8320; (layer <span id="pl"></span>) across tokens</h2><div id="clock"></div></div>
<div class="panel"><h2>Hamiltonian constraint |H| across tokens (lower = more physical, H&Psi;=0)</h2><div id="hcon"></div></div>
<div class="panel"><h2>Minisuperspace state &mdash; K modes &Psi; at this token (orange = timelike)</h2><div id="modes"></div></div>
</div>
<div style="display:flex;flex-direction:column;gap:12px">
<div class="panel"><h2>|H| per layer (this token)</h2><div id="heat"></div></div>
<div class="panel"><h2>Attention &mdash; layer <span id="al"></span> (last token &rarr; context)</h2><div id="attn"></div></div>
<div class="panel"><h2>Notes</h2><div class="small" id="notes"></div></div>
</div>
</div>
<script>
const D=/*__DATA__*/; const M=D.meta,F=D.frames,NL=M.n_layer,K=M.wdw_modes,PL=M.phase_layer;
let i=0,playing=false,timer=null; const $=id=>document.getElementById(id);
$("meta").textContent=`${M.ckpt} · step ${M.step} · ${K} modes · ${M.wdw_steps} wave steps · ${NL} layers · "${M.prompt}"`;
$("pl").textContent=PL; $("al").textContent=PL;
const seek=$("seek"); seek.max=F.length-1;
const hue2=v=>`hsl(${(1-Math.max(0,Math.min(1,v)))*20+0},80%,${30+35*Math.max(0,Math.min(1,v))}%)`; // red hot heat
// per-token series from the phase layer: Psi_0 (clock) and |H| (constraint residual)
const clock=F.map(f=>f.rings&&f.rings[PL]?f.rings[PL].psi[0]:0);
const hcon =F.map(f=>f.rings&&f.rings[PL]?Math.abs(f.rings[PL].R[0]):0);
// generic SVG line chart of a series, with a cursor at index i and a zero baseline
function lineChart(series,color,W,H,cur,zero){
const n=series.length; if(!n) return "";
let mn=Math.min(...series),mx=Math.max(...series); if(zero){mn=Math.min(mn,0);mx=Math.max(mx,0);}
if(mx-mn<1e-9)mx=mn+1; const pad=6;
const X=k=>pad+(W-2*pad)*(n<2?0.5:k/(n-1));
const Y=v=>pad+(H-2*pad)*(1-(v-mn)/(mx-mn));
let pts=series.map((v,k)=>`${X(k).toFixed(1)},${Y(v).toFixed(1)}`).join(" ");
let z = zero?`<line x1="${pad}" y1="${Y(0).toFixed(1)}" x2="${W-pad}" y2="${Y(0).toFixed(1)}" stroke="#30363d" stroke-dasharray="3 3"/>`:"";
let cx=X(cur).toFixed(1);
let cy=Y(series[cur]).toFixed(1);
return `<svg viewBox="0 0 ${W} ${H}">${z}
<polyline points="${pts}" fill="none" stroke="${color}" stroke-width="1.6"/>
<line x1="${cx}" y1="${pad}" x2="${cx}" y2="${H-pad}" stroke="${color}" stroke-opacity="0.4"/>
<circle cx="${cx}" cy="${cy}" r="3.2" fill="${color}"/>
<text x="${pad}" y="12" fill="#8b949e" font-size="10">${mx.toFixed(3)}</text>
<text x="${pad}" y="${H-3}" fill="#8b949e" font-size="10">${mn.toFixed(3)}</text></svg>`;
}
// K minisuperspace modes as up/down bars around a zero baseline; mode 0 (timelike) hot
function modeBars(theta){
const n=theta.length,W=560,H=140,pad=8; if(!n)return "";
const amx=Math.max(1e-6,...theta.map(Math.abs));
const bw=(W-2*pad)/n, mid=H/2;
let bars=theta.map((v,k)=>{const hgt=(H/2-pad)*(Math.abs(v)/amx);
const y=v>=0?mid-hgt:mid; const col=(k===0)?"var(--hot)":"var(--ac)";
return `<rect x="${(pad+k*bw).toFixed(1)}" y="${y.toFixed(1)}" width="${Math.max(1,bw-1).toFixed(1)}" height="${hgt.toFixed(1)}" fill="${col}"/>`;}).join("");
return `<svg viewBox="0 0 ${W} ${H}"><line x1="${pad}" y1="${mid}" x2="${W-pad}" y2="${mid}" stroke="#30363d"/>${bars}
<text x="${pad}" y="12" fill="#8b949e" font-size="10">timelike &Psi;&#8320; = ${theta[0].toFixed(3)}</text></svg>`;
}
const stream=$("stream");
D.prompt_chars.forEach(c=>{const s=document.createElement("span");s.className="tk prompt";s.textContent=esc(c);stream.appendChild(s);});
F.forEach((f,k)=>{const s=document.createElement("span");s.className="tk gen";s.textContent=esc(f.char);s.onclick=()=>{i=k;render()};s.dataset.k=k;stream.appendChild(s);});
function esc(c){return c.replace(/\n/g,"⏎").replace(/ /g,"·");}
function render(){
const f=F[i]; $("pos").textContent=`${i+1}/${F.length}`; seek.value=i;
document.querySelectorAll(".tk.gen").forEach(s=>s.classList.toggle("cur",+s.dataset.k===i));
$("clock").innerHTML=lineChart(clock,"#f0883e",560,150,i,true);
$("hcon").innerHTML =lineChart(hcon, "#f85149",560,110,i,false);
$("modes").innerHTML=(f.phases&&f.phases.theta)?modeBars(f.phases.theta):"<span class='small'>n/a</span>";
// |H| per layer heat row
const hL=(f.rings||[]).map(r=>Math.abs(r.R[0])); const hmx=Math.max(...hL,1e-6);
let hh=`<div class="grid" style="grid-template-columns:repeat(${NL},1fr)">`;
for(let l=0;l<NL;l++){const v=hL[l]||0;hh+=`<div class="cell" title="L${l} |H|=${v.toFixed(4)}" style="background:${hue2(v/hmx)}"></div>`;}
$("heat").innerHTML=hh+"</div><div class='small'>bright red = far from H=0</div>";
// trait activity (block norms + HRM/MoE)
const qn=f.quaz_norm||[],mx=Math.max(...qn,1e-6);let L=qn.map((_,l)=>"wave L"+l),V=qn.map(v=>v/mx);
if(f.traits&&f.traits.hrm!=null){L.push("HRM");V.push(Math.min(1,f.traits.hrm/mx));}
if(f.traits&&f.traits.moe!=null){L.push("MoE");V.push(Math.min(1,f.traits.moe/mx));}
let th="<table>";V.forEach((v,k)=>{th+=`<tr><td class="lbl">${L[k]}</td><td style="width:100%"><div class="bar"><div style="width:${(v*100).toFixed(0)}%;background:${hue2(v)}"></div></div></td></tr>`;});
$("traits").innerHTML=th+"</table>";
let tp="<table>";f.top.forEach(([t,p])=>{tp+=`<tr><td class="lbl">${p.toFixed(2)}</td><td><div class="bar"><div style="width:${(p*100).toFixed(0)}%;background:var(--ac)"></div></div></td><td class="small">id ${t}</td></tr>`;});
$("top").innerHTML=tp+"</table>";
if(f.attn){const w=f.attn.w,m=Math.max(...w,1e-6);let h="<div style='display:flex;flex-wrap:wrap;gap:1px'>";
w.forEach((a,p)=>{h+=`<div title="pos ${p}: ${a.toFixed(3)}" style="width:8px;height:14px;background:${hue2(a/m)}"></div>`;});
$("attn").innerHTML=h+`</div><div class='small'>${w.length} context positions</div>`;}
else $("attn").innerHTML="<span class='small'>n/a</span>";
$("notes").innerHTML=`Modes: ${K} · wave steps: ${M.wdw_steps}<br>&Psi;&#8320; is the volume/scale (timelike) mode &mdash; the emergent clock; there is no external time.<br>|H| is the Hamiltonian-constraint residual; training drives it toward 0 (H&Psi;=0, a physical state).<br>Orange bar = timelike mode; blue = spacelike. Click any token or use the slider.`;
}
seek.oninput=()=>{i=+seek.value;render()};
$("play").onclick=()=>{playing=!playing;$("play").textContent=playing?"⏸ pause":"▶ play";if(playing)timer=setInterval(()=>{i=(i+1)%F.length;render();},400);else clearInterval(timer);};
document.onkeydown=e=>{if(e.key==="ArrowRight"){i=Math.min(F.length-1,i+1);render();}if(e.key==="ArrowLeft"){i=Math.max(0,i-1);render();}};
render();
</script></body></html>"""
if __name__ == "__main__":
main()