Wheeler-Chat / app.py
Quazim0t0's picture
Fix infinite page growth: remove fill_height (feedback loop with HF iframe auto-resizer)
ddc5034
Raw
History Blame Contribute Delete
24.1 kB
"""
Wheeler Chat + Dynamics -- a Gradio demo for Wheeler-DeWitt-62M, a language model whose channel
mixer is the Wheeler-DeWitt equation of quantum gravity. You chat with the SFT checkpoint
AND watch the Wheeler-DeWitt dynamics that produced each reply: the emergent-time mode
Psi_0 (the "clock"), the Hamiltonian-constraint residual |H|, and the K minisuperspace
modes, captured token-by-token during generation and rendered as an inline-SVG dashboard.
Weights are pulled at runtime from the public repo Quazim0t0/Wheeler-DeWitt-62M. Research artifact: the SFT checkpoint learned the FORM of a chat turn but not
knowledge -- it hallucinates and does not reliably follow instructions.
"""
import re, os, json, html, sys
try: sys.stdout.reconfigure(line_buffering=True)
except Exception: pass
print(">> app: importing torch", flush=True)
import torch
import torch.nn.functional as F
print(">> app: torch ok, importing gradio", flush=True)
import gradio as gr
from huggingface_hub import hf_hub_download
print(">> app: gradio ok, importing model chain", flush=True)
from model import QuazimotoLM, QuazimotoConfig
from spike_tokenizer import SpikeTokenizer
import instrument
from visualize import HTML_TEMPLATE
print(">> app: all imports done, building UI", flush=True) # the inline-SVG Wheeler dashboard
REPO = "Quazim0t0/Wheeler-DeWitt-62M"
TOKEN = os.environ.get("HF_TOKEN")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") # faster 755MB pull
ROLE = {"system": "<|system|>", "user": "<|user|>", "assistant": "<|assistant|>"}
# Lazy load: bind the UI first (so the Space reaches "Running"), then load the
# 62.9M checkpoint on the first request. Loading at import time blocks the port
# past the startup health-check -> the Space hangs on "starting".
model = tok = cfg = CK = VOCAB = IM_END = STOP_IDS = PL = None
def _ensure_loaded():
global model, tok, cfg, CK, VOCAB, IM_END, STOP_IDS, PL
if model is not None:
return
ckpt_path = hf_hub_download(REPO, "wheeler-dewitt-62m-sft.pt", token=TOKEN)
tok = SpikeTokenizer(vocab_file="tokenizer.json")
CK = torch.load(ckpt_path, map_location="cpu", weights_only=False)
cfg = QuazimotoConfig(**CK["family_config"])
model = QuazimotoLM(cfg); model.load_state_dict(CK["model"], strict=False); model.eval()
VOCAB = tok.get_vocab()
IM_END = VOCAB.get("<|im_end|>")
STOP_IDS = {VOCAB[n] for n in ("<|im_end|>", "<eos>", "<|endoftext|>") if n in VOCAB}
PL = cfg.n_layer // 2
def render(history, system, user):
parts = []
if system:
parts.append(f"<|im_start|>{ROLE['system']}\n{system}<|im_end|>\n")
for m in (history or []):
if not isinstance(m, dict): # tolerate legacy tuple entries
m = {"role": "assistant", "content": m[1] if len(m) > 1 else ""}
r, c = m.get("role"), _as_text(m.get("content", ""))
if r in ("user", "assistant") and c:
parts.append(f"<|im_start|>{ROLE[r]}\n{c}<|im_end|>\n")
parts.append(f"<|im_start|>{ROLE['user']}\n{user}<|im_end|>\n<|im_start|>{ROLE['assistant']}\n")
return "".join(parts)
def _ban3(ids, n=3):
"""no-repeat-ngram: token ids that would complete an already-seen n-gram.
A/B on real weights (5 prompts, 120 tok): rep-4 0.181 -> 0.000, distinct-2
0.745 -> 0.913 at the deployed sampling settings."""
if len(ids) < n:
return ()
tail = tuple(ids[-(n - 1):])
return tuple(ids[i + n - 1] for i in range(len(ids) - n + 1)
if tuple(ids[i:i + n - 1]) == tail)
@torch.no_grad()
def _one_candidate(history, system, user, temperature, top_p, max_new, seed=0):
"""Generate the reply while recording the per-token Wheeler-DeWitt state (via
instrument.Recorder, which WheelerDeWittBlock writes to). Returns (text, frames).
KV-cached like generate.py: the prompt is prefilled once, then each new token
is a single-position forward (a full recompute per token is O(n^2) -- minutes
on CPU). The recorder captures the LAST position of every forward, so each
single-token step still yields exactly one clean frame for the token being
processed. The first token is sampled from the prefill logits."""
_ensure_loaded() # first call loads the model (lazy)
torch.manual_seed(1234 + seed) # reproducible per-candidate sampling
if hasattr(model, "reset_ring_memory"):
model.reset_ring_memory() # clean per-request fast-weight memory
rec = instrument.Recorder(phase_layer=PL, attn_layer=PL)
ids = tok.encode(render(history, system, user), add_special_tokens=False)
ids = ids[-(cfg.block_size - int(max_new) - 4):]
idx = torch.tensor([ids]); out = []
def sample(lg1):
lg = torch.nan_to_num(lg1.float(), nan=0.0, posinf=1e4, neginf=-1e4)
if out: # repetition penalty 1.15 (reply only) -- nrng
seen = torch.tensor(sorted(set(out)))
v = lg[0, seen]; lg[0, seen] = torch.where(v > 0, v / 1.15, v * 1.15)
for b in _ban3(out):
lg[0, b] = -float("inf")
lg = lg / max(float(temperature), 1e-6)
k = min(40, lg.size(-1)) # top-k 40 -- nrng
lg = lg.masked_fill(lg < torch.topk(lg, k).values[..., -1, None], -float("inf"))
probs = F.softmax(lg, dim=-1)
if float(top_p) < 1.0:
sp, si = probs.sort(-1, descending=True)
cut = sp.cumsum(-1) > float(top_p)
cut[..., 1:] = cut[..., :-1].clone(); cut[..., 0] = False
sp = sp.masked_fill(cut, 0.0); sp = sp / sp.sum(-1, keepdim=True)
return si.gather(-1, torch.multinomial(sp, 1))
return torch.multinomial(probs, 1)
try:
logits, _, _, past = model(idx, use_cache=True) # prefill (no capture)
nxt = sample(logits[:, -1, :])
for _ in range(int(max_new)):
t = int(nxt)
if t in STOP_IDS:
break
out.append(t); idx = torch.cat([idx, nxt], dim=1)
# capture the forward that PROCESSES this token -> one frame per token
instrument.set_rec(rec); rec.begin()
if past[0][0].size(2) < cfg.block_size:
logits, _, _, past = model(idx[:, -1:], past_key_values=past, use_cache=True)
else: # context full -> windowed re-prime
logits, _, _, past = model(idx[:, -cfg.block_size:], use_cache=True)
topk = torch.topk(F.softmax(logits[:, -1].float(), -1), 5)
rec.end(token=t, char=tok.decode([t], skip_special_tokens=False),
top=[[int(i), round(float(p), 3)] for i, p in zip(topk.indices[0], topk.values[0])])
instrument.set_rec(None)
nxt = sample(logits[:, -1, :])
finally:
instrument.set_rec(None)
return tok.decode(out, skip_special_tokens=True), rec.frames
# ---------------------------------------------------------------------------
# Effort-scaled decoding (candidate sampling + scoring) is adapted from Glint-2
# (https://huggingface.co/Glint-Research/Glint-2), Copyright (c) Glint-Research,
# MIT License.
# ---------------------------------------------------------------------------
EFFORT_LEVELS = {"fast (1 sample)": 1, "balanced (3 samples)": 3, "best (6 samples)": 6}
_ETEMPS = [0.30, 0.35, 0.40, 0.45]
_STOPW = {"the","a","an","is","are","do","does","how","what","why","of","in","to",
"for","and","or","it","its","with","on","at","by","that","this","when",
"from","i","you","my","your","can","should","was","were","be","have"}
def _as_text(x):
"""Coerce a history entry / content field to plain text.
This Space's Chatbot runs in Gradio's TUPLES format while on_submit appends
{"role","content"} dicts, so history arrives mixed: some entries are
[user, assistant] lists, some are dicts, and a dict's "content" can itself be
a list of parts. Be liberal rather than assume one shape."""
if x is None:
return ""
if isinstance(x, str):
return x
if isinstance(x, dict):
# gradio message dicts use "content"; multimodal content PARTS use "text"
for k in ("content", "text", "value", "alt_text"):
if x.get(k) is not None:
return _as_text(x[k])
return ""
if isinstance(x, (list, tuple)):
return " ".join(_as_text(i) for i in x)
return str(x)
def _content(t):
return {w for w in re.findall(r"[a-z]+", _as_text(t).lower())
if w not in _STOPW and len(w) > 3}
def _score_reply(text, user, prev):
"""Rank candidates: penalise self-repetition, reward addressing the CURRENT
question, penalise echoing the previous assistant turn."""
words = re.findall(r"[a-z]+", _as_text(text).lower())
if not words:
return -1e9
grams = [tuple(words[i:i + 4]) for i in range(len(words) - 3)]
rep = 1.0 - len(set(grams)) / len(grams) if grams else 0.0
w = set(words)
s = -2.0 * rep + 0.02 * min(len(words), 60)
s += 0.6 * min(len(_content(user) & w), 4)
s -= 0.5 * min(len(_content(prev) & w), 4)
return s
def generate_with_capture(history, system, user, effort, top_p, max_new,
progress_cb=None):
n = EFFORT_LEVELS.get(effort, 3)
prev = ""
for h in reversed(history or []):
if isinstance(h, dict):
if h.get("role") == "assistant" and h.get("content"):
prev = _as_text(h["content"]); break
elif isinstance(h, (list, tuple)) and len(h) > 1 and h[1]:
prev = _as_text(h[1]); break # tuples format: [user, assistant]
best, best_frames, best_s = "", [], -1e18
for i in range(n):
if progress_cb is not None:
progress_cb(i, n)
temp = 0.7 if n == 1 else _ETEMPS[i % len(_ETEMPS)]
text, frames = _one_candidate(history, system, user, temp, top_p, max_new, seed=i)
sc = _score_reply(text, user, prev)
if sc > best_s:
best, best_frames, best_s = text, frames, sc
return best, best_frames
def dashboard_html(frames, prompt):
"""Embed the visualize.py inline-SVG dashboard (its JS runs in a sandboxed iframe)."""
if not frames:
return "<div style='color:#9a94c0;padding:22px;text-align:center;font-size:.75rem'>No dynamics captured.</div>"
data = {"meta": {"ckpt": "Wheeler-DeWitt-62M (SFT)", "step": CK.get("step"),
"n_layer": cfg.n_layer, "wdw_modes": cfg.wdw_modes,
"wdw_steps": cfg.wdw_steps, "phase_layer": PL, "prompt": prompt[:60]},
"prompt_chars": [], "frames": frames}
page = HTML_TEMPLATE.replace("/*__DATA__*/", json.dumps(data))
return (f'<iframe srcdoc="{html.escape(page, quote=True)}" '
f'style="width:100%;height:660px;border:0;border-radius:8px"></iframe>')
def on_submit(user, history, system, effort, top_p, max_new,
progress=gr.Progress()):
"""Ordinary (non-generator) handler. Progress is reported through gradio's
Progress bar rather than by yielding chat messages: yielding placeholders
into the transcript made them accumulate in the UI instead of replacing.
The progress call is wrapped so any API difference degrades to no bar
rather than breaking generation."""
history = (history or []) + [{"role": "user", "content": user}]
def _tick(i, n):
try:
progress((i) / max(n, 1),
desc="searching" if n == 1 else f"searching {i + 1}/{n}")
except Exception:
pass
reply, frames = generate_with_capture(history[:-1], system, user, effort,
top_p, max_new, progress_cb=_tick)
history = history + [{"role": "assistant", "content": reply or "…"}]
return history, dashboard_html(frames, user), ""
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@200;300;400;500;600&family=JetBrains+Mono:wght@400;500&display=swap');
:root { --bg:#04030b; --panel:rgba(160,140,255,.04); --panel2:rgba(160,140,255,.06);
--txt:#ece9fb; --muted:#9a94c0; --dim:#5c5680; --border:rgba(160,140,255,.12);
--g1:#8b6bff; --g2:#ff5fd0; --g3:#53d6ff;
--spectrum:linear-gradient(90deg,#53d6ff,#8b6bff,#ff5fd0,#8b6bff,#53d6ff); }
gradio-app { background:transparent !important; }
body { background:radial-gradient(1400px 900px at 50% -20%, #120b2e 0%, #04030b 60%) fixed !important; }
.gradio-container { background:transparent !important; color:var(--txt) !important;
font-family:'Inter','SF Pro Display',ui-sans-serif,-apple-system,'Segoe UI',sans-serif !important;
max-width:1240px !important; margin:0 auto !important; position:relative; }
.gradio-container::before { content:''; position:fixed; inset:0; z-index:0; pointer-events:none;
background:radial-gradient(ellipse 120% 90% at 50% 40%, transparent 55%, rgba(0,0,0,.6) 100%); }
.gradio-container::after { content:''; position:fixed; inset:0; z-index:0; pointer-events:none; opacity:.04;
background-image:url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='160' height='160'%3E%3Cfilter id='n'%3E%3CfeTurbulence type='fractalNoise' baseFrequency='0.9' numOctaves='2'/%3E%3C/filter%3E%3Crect width='100%25' height='100%25' filter='url(%23n)'/%3E%3C/svg%3E");
}
.gradio-container > * { position:relative; z-index:1; }
/* hero */
#hero { text-align:center; padding:30px 0 4px; }
#hero .tag { display:inline-block; font-family:'JetBrains Mono',monospace; font-size:.62rem;
letter-spacing:.3em; text-transform:uppercase; color:var(--muted);
border:1px solid var(--border); border-radius:999px; padding:5px 16px; margin-bottom:14px;
background:rgba(160,140,255,.04); animation:rise .8s cubic-bezier(.16,1,.3,1) both; }
#hero .tag b { color:var(--g3); font-weight:500; }
#hero h1 { font-size:clamp(2.4rem,6vw,3.8rem); font-weight:200; margin:0; letter-spacing:-.015em;
line-height:1.02; color:var(--txt); animation:rise .9s cubic-bezier(.16,1,.3,1) .08s both; }
#hero h1 em { font-style:normal; background-image:var(--spectrum) !important;
background-size:200% 100% !important; -webkit-background-clip:text !important;
background-clip:text !important; -webkit-text-fill-color:transparent;
}
@keyframes shimmer { to { background-position:200% 0; } }
#hero .rule { width:150px; height:1px; margin:16px auto 0; background:var(--spectrum);
background-size:200% 100%; opacity:.85; }
#hero p { color:var(--muted); margin:.7rem auto 0; font-size:.88rem; font-weight:300; max-width:72ch;
animation:rise 1s cubic-bezier(.16,1,.3,1) .16s both; }
#hero .warn { display:inline-block; margin-top:.7rem; font-family:'JetBrains Mono',monospace;
font-size:.66rem; letter-spacing:.06em; color:#d9a15c; border:1px solid rgba(217,161,92,.25);
background:rgba(217,161,92,.06); border-radius:8px; padding:6px 14px; max-width:80ch;
animation:rise 1.1s cubic-bezier(.16,1,.3,1) .24s both; }
@keyframes rise { from { opacity:0; transform:translateY(16px); } to { opacity:1; transform:none; } }
/* glass, compact */
.gradio-container .block { padding:8px 10px !important; }
.gradio-container .form { padding:6px !important; gap:6px !important; }
.gradio-container .gap, .gradio-container .row, .gradio-container .column { gap:10px !important; }
.block, .form, .accordion, .gr-group { background:var(--panel) !important;
border:1px solid var(--border) !important; border-radius:16px !important; }
div[data-testid='chatbot'], .chatbot, div[data-testid='chatbot'] > div,
.bubble-wrap, .message-wrap { background:transparent !important; border:none !important; }
[data-testid='block-label'], label.float { display:none !important; }
.message, .bubble { border-radius:14px !important; border:1px solid var(--border) !important;
background:var(--panel2) !important; color:var(--txt) !important;
padding:8px 13px !important; font-size:.88rem !important; line-height:1.45 !important; }
.user .message, .user .bubble { background:linear-gradient(92deg,rgba(139,107,255,.16),rgba(255,95,208,.14)) !important; }
button.primary, .gr-button-primary { background:linear-gradient(92deg,var(--g1),var(--g2)) !important;
border:none !important; color:#fff !important; font-weight:500 !important; border-radius:12px !important;
padding:7px 16px !important; transition:filter .2s, box-shadow .2s; }
button.primary:hover { filter:brightness(1.12); box-shadow:0 0 26px rgba(139,107,255,.4); }
button.secondary { background:var(--panel) !important; border:1px solid var(--border) !important;
color:var(--muted) !important; border-radius:12px !important; }
textarea, input[type=text], input[type=number] { background:rgba(160,140,255,.05) !important;
color:var(--txt) !important; border:1px solid var(--border) !important; border-radius:12px !important;
font-size:.88rem !important; }
textarea:focus, input[type=text]:focus { border-color:rgba(139,107,255,.55) !important;
box-shadow:0 0 0 3px rgba(139,107,255,.14) !important; }
label, .label-wrap span, span[data-testid='block-info'] { color:var(--muted) !important;
font-family:'JetBrains Mono',monospace !important; font-size:.7rem !important;
letter-spacing:.08em; text-transform:uppercase; background:transparent !important; padding:0 !important; }
.examples button, button.example { background:var(--panel) !important; border:1px solid var(--border) !important;
border-radius:999px !important; color:var(--muted) !important; font-size:.75rem !important;
padding:5px 12px !important; transition:all .2s; }
.examples button:hover { border-color:rgba(83,214,255,.5) !important; color:var(--txt) !important; }
.hero-wrap, .hero-wrap.block, .block:has(#hero) { background:transparent !important;
border:none !important; backdrop-filter:none !important; box-shadow:none !important; }
footer { display:none !important; }
"""
HEAD = """
<script>
function __bgScene() {
// Static scene: rendered exactly once. No animation loop, so it can never
// flicker, strobe, or consume GPU after load.
const cv = document.createElement('canvas');
cv.style.cssText = 'position:fixed;inset:0;width:100vw;height:100vh;z-index:0;pointer-events:none;';
document.body.prepend(cv);
const ctx = cv.getContext('2d');
const W = cv.width = innerWidth, H = cv.height = innerHeight;
const rnd = (a, b) => a + Math.random() * (b - a);
const TAU = Math.PI * 2;
ctx.globalCompositeOperation = 'lighter';
// nebula
const blobs = [[0.22,0.20,0.36,'139,107,255',0.10],[0.78,0.16,0.32,'255,95,208',0.07],
[0.55,0.75,0.42,'83,214,255',0.05],[0.40,0.45,0.30,'104,60,220',0.08]];
const M = Math.min(W, H);
for (const b of blobs) {
const g = ctx.createRadialGradient(b[0]*W, b[1]*H, 0, b[0]*W, b[1]*H, b[2]*M);
g.addColorStop(0, `rgba(${b[3]},${b[4]})`); g.addColorStop(1, `rgba(${b[3]},0)`);
ctx.fillStyle = g; ctx.fillRect(0, 0, W, H);
}
// stars
for (let i = 0; i < 170; i++) {
const z = rnd(0.3, 1), r = rnd(0.4, 1.3) * z;
ctx.fillStyle = `rgba(226,225,248,${rnd(0.2, 0.8) * z})`;
ctx.fillRect(rnd(0, W), rnd(0, H), r, r);
}
// the orrery, centered under the hero
const CX = W * 0.5, CY = H * 0.34;
const COLS = ['83,214,255', '139,107,255', '255,95,208'];
for (let i = 0; i < 5; i++) {
const R = (0.10 + i * 0.052) * M * 1.35, tilt = rnd(0.30, 0.62), rot = rnd(0, TAU);
const c = COLS[i % 3];
ctx.save(); ctx.translate(CX, CY); ctx.rotate(rot);
ctx.strokeStyle = `rgba(${c},0.16)`; ctx.lineWidth = 1;
ctx.beginPath(); ctx.ellipse(0, 0, R, R * tilt, 0, 0, TAU); ctx.stroke();
for (let q = 0; q < 1 + (i % 3); q++) {
const th = rnd(0, TAU);
const qx = Math.cos(th) * R, qy = Math.sin(th) * R * tilt;
const g = ctx.createRadialGradient(qx, qy, 0, qx, qy, 9);
g.addColorStop(0, `rgba(${c},0.9)`); g.addColorStop(1, `rgba(${c},0)`);
ctx.fillStyle = g; ctx.beginPath(); ctx.arc(qx, qy, 9, 0, TAU); ctx.fill();
ctx.fillStyle = 'rgba(255,255,255,0.95)';
ctx.fillRect(qx - 1.5, qy - 1.5, 3, 3);
}
ctx.restore();
}
// central glow
const cg = ctx.createRadialGradient(CX, CY, 0, CX, CY, M * 0.06);
cg.addColorStop(0, 'rgba(236,233,251,0.18)'); cg.addColorStop(1, 'rgba(139,107,255,0)');
ctx.fillStyle = cg; ctx.beginPath(); ctx.arc(CX, CY, M * 0.06, 0, TAU); ctx.fill();
}
if (document.readyState === 'loading') document.addEventListener('DOMContentLoaded', __bgScene);
else __bgScene();
</script>
"""
HERO = (
"<div id='hero'>"
"<span class='tag'><b>Wheeler-DeWitt-62M</b> &nbsp;&middot;&nbsp; quantum-gravity LM "
"&nbsp;&middot;&nbsp; live dynamics</span>"
"<h1>The <em>Wheeler</em> Universe</h1><div class='rule'></div>"
"<p>A 62.9M-parameter language model whose channel mixer is built from the "
"<b>Wheeler&ndash;DeWitt equation</b> of quantum gravity &mdash; a <i>timeless</i> constraint "
"with a Lorentzian supermetric &mdash; in place of the usual MLP. Chat below; the dynamics "
"panel shows the physics that produced each reply: the emergent-time clock &Psi;&#8320;, the "
"Hamiltonian constraint |H|, and the minisuperspace modes, captured token-by-token.</p>"
"<span class='warn'>&#9888; research artifact &mdash; WheelerV2 tends to follow conversations "
"well and understands somewhat, but is not trained on enough data to be reliable on many facts</span>"
"</div>"
)
with gr.Blocks(title="Wheeler Chat", css=CSS, head=HEAD) as demo:
gr.HTML(elem_classes=["hero-wrap"], value=HERO)
with gr.Row():
with gr.Column(scale=1):
# NOTE: do not pass type= here. gradio 6 removed the parameter
# (messages is the only format); gradio 5 warns without it. The
# Space pins sdk_version 6.15.2, so omit it.
chatbot = gr.Chatbot(height=380, label="Wheeler-DeWitt-62M")
msg = gr.Textbox(placeholder="Ask Wheeler-DeWitt-62M anything…",
label="Message")
with gr.Accordion("settings", open=False):
system = gr.Textbox(value="You are a helpful assistant.", label="System prompt")
temperature = gr.Dropdown(choices=list(EFFORT_LEVELS.keys()),
value="balanced (3 samples)", label="effort")
top_p = gr.Slider(0.1, 1.0, value=0.95, label="top_p")
max_new = gr.Slider(16, 160, value=64, step=8, label="max new tokens")
with gr.Row():
send = gr.Button("Send", variant="primary")
clear = gr.Button("Clear")
gr.Examples(
examples=[["What is the capital of France?"],
["Explain photosynthesis in one sentence."],
["Tell me a short story."],
["Who described the three laws of motion?"],
["How do you bake a loaf of bread?"]],
inputs=msg, label="Try")
viz = gr.HTML(label="Wheeler–DeWitt dynamics",
value="<div style='color:#9a94c0;padding:22px;text-align:center;font-size:.75rem;letter-spacing:.08em'>Send a message to watch "
"the emergent-time clock and the Hamiltonian constraint evolve.</div>")
inp = [msg, chatbot, system, temperature, top_p, max_new]
send.click(on_submit, inp, [chatbot, viz, msg])
msg.submit(on_submit, inp, [chatbot, viz, msg])
clear.click(lambda: ([], "", ""), None, [chatbot, viz, msg])
print(">> app: launching gradio", flush=True)
if __name__ == "__main__":
demo.queue().launch()