app.py

#!/usr/bin/env python3
"""
app.py — CDM V6 HORN Demo — HuggingFace Spaces
DuoNeural (Archon + Jesse + Aura) — 2026

Shows CDM V6 HORN in action:
  - Story generation with slot dynamics visualization
  - Oscillator regime display (γ/ω per layer, underdamped vs overdamped)
  - Slot logit lens: what each memory slot is tracking at each generation step
"""

import json
import math
import torch
import gradio as gr
import torch.nn.functional as F
from huggingface_hub import hf_hub_download

model       = None
tokenizer   = None
osc_html    = None   # cached oscillator panel HTML
DEVICE      = "cpu"  # HF Spaces CPU tier

SLOT_COLORS = [
    "#FF6B6B","#4ECDC4","#45B7D1","#96CEB4",
    "#FECA57","#FF9FF3","#54A0FF","#5F27CD",
    "#00D2D3","#FF9F43","#C8D6E5","#8395A7",
    "#EE5A24","#009432","#C4E538","#A3CB38",
]

REGIME_COLORS = {
    "underdamped": "#4ECDC4",
    "overdamped":  "#FF9F43",
    "critical":    "#FECA57",
}

EXAMPLE_PROMPTS = [
    "Once upon a time there was a little girl named Lily.",
    "Tom loved trains more than anything else in the world.",
    "The rabbit hopped through the sunny meadow looking for",
    "In a small village there lived a clever fox named",
    "Sara and Ben decided to build the tallest tower ever",
]


# ─── Model loading ────────────────────────────────────────────────────────────

def load_model():
    global model, tokenizer, osc_html
    if model is not None:
        return

    from transformers import GPT2TokenizerFast
    from cdm_model_v6_horn import CDMLanguageModelV6HORN
    from cdm_model_v3 import CDMConfigV3

    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token

    model_path = hf_hub_download(
        repo_id="DuoNeural/CDM-V6-HORN-TinyStories-37M",
        filename="model.pt"
    )
    cfg_path = hf_hub_download(
        repo_id="DuoNeural/CDM-V6-HORN-TinyStories-37M",
        filename="config.json"
    )

    with open(cfg_path) as f:
        cfg_dict = json.load(f)

    cfg = CDMConfigV3(
        vocab_size  = cfg_dict.get("vocab_size", 50257),
        d_model     = cfg_dict.get("d_model", 384),
        n_layers    = cfg_dict.get("n_layers", 8),
        n_heads     = cfg_dict.get("n_heads", 8),
        n_kv_heads  = cfg_dict.get("n_kv_heads", 4),
        d_ff        = cfg_dict.get("d_ff", 1024),
        K           = cfg_dict.get("K", 16),
        max_len     = cfg_dict.get("max_len", 512),
        lbl_coeff   = cfg_dict.get("lbl_coeff", 0.01),
        entropy_reg = cfg_dict.get("entropy_reg", 0.02),
    )

    model_obj = CDMLanguageModelV6HORN(cfg)
    ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
    state = ckpt.get("model_state", ckpt)
    model_obj.load_state_dict(state, strict=False)
    model_obj.eval()
    model = model_obj
    osc_html = _build_oscillator_panel()


# ─── Oscillator panel (static, built once after model load) ───────────────────

def _build_oscillator_panel():
    rows = []
    for l_idx, block in enumerate(model.blocks):
        with torch.no_grad():
            gamma = F.softplus(block.cdm.raw_gamma).numpy()
            omega = F.softplus(block.cdm.raw_omega).numpy()

        g_mean = float(gamma.mean())
        o_mean = float(omega.mean())
        under_frac = float((omega > gamma).mean())
        tau_L = 1.0 / g_mean
        tau_star = 0.72 / g_mean

        if under_frac > 0.6:
            regime = "underdamped"
            regime_label = f"UNDERDAMPED ({under_frac:.0%}) — rings"
        elif under_frac < 0.3:
            regime = "overdamped"
            regime_label = f"OVERDAMPED ({1-under_frac:.0%}) — stable"
        else:
            regime = "critical"
            regime_label = f"MIXED ({under_frac:.0%} under)"

        color = REGIME_COLORS[regime]

        bars = "▁▂▃▄▅▆▇█"
        g_min, g_max = gamma.min(), gamma.max()
        if g_max > g_min:
            sparkline = "".join(bars[int((v - g_min) / (g_max - g_min) * 7)] for v in gamma)
        else:
            sparkline = "─" * len(gamma)

        rows.append(f"""
<tr>
  <td style="color:#aaa;padding:4px 8px;font-weight:bold;">L{l_idx}</td>
  <td style="color:#ff9ff3;padding:4px 8px;font-family:monospace;">γ={g_mean:.3f}</td>
  <td style="color:#4ecdc4;padding:4px 8px;font-family:monospace;">ω={o_mean:.3f}</td>
  <td style="color:#feca57;padding:4px 8px;font-family:monospace;">τ_L={tau_L:.2f} | τ*={tau_star:.2f}</td>
  <td style="color:{color};padding:4px 8px;font-size:11px;">{regime_label}</td>
  <td style="color:#666;padding:4px 8px;font-family:monospace;font-size:10px;">{sparkline}</td>
</tr>""")

    return f"""
<div style="background:#0d0d1a;padding:12px;border-radius:8px;border:1px solid #333;">
<p style="color:#888;font-size:12px;margin:0 0 8px 0;">
  <b style="color:#4ecdc4;">Underdamped</b> (ω&gt;γ): slot rings when written → reactive, resonant<br>
  <b style="color:#ff9f43;">Overdamped</b>  (γ&gt;ω): slot returns smoothly → stable persistent storage<br>
  <b style="color:#aaa;">τ_L</b>=1/γ (relaxation time) | <b style="color:#aaa;">τ*</b>=0.72/γ (DHP predictability horizon)
</p>
<table style="border-collapse:collapse;width:100%;font-size:12px;">
  <thead>
    <tr style="border-bottom:1px solid #333;">
      <th style="color:#fff;padding:4px 8px;text-align:left;">Layer</th>
      <th style="color:#ff9ff3;padding:4px 8px;text-align:left;">Damping γ</th>
      <th style="color:#4ecdc4;padding:4px 8px;text-align:left;">Freq ω</th>
      <th style="color:#feca57;padding:4px 8px;text-align:left;">Timescales</th>
      <th style="color:#fff;padding:4px 8px;text-align:left;">Regime</th>
      <th style="color:#666;padding:4px 8px;text-align:left;">Per-slot γ</th>
    </tr>
  </thead>
  <tbody>{"".join(rows)}</tbody>
</table>
</div>"""


# ─── Generation ───────────────────────────────────────────────────────────────

@torch.no_grad()
def generate_with_slots(prompt: str, max_new_tokens: int, temperature: float, top_k: int):
    load_model()

    ids = tokenizer.encode(prompt, return_tensors="pt")   # (1, T_prompt)
    K = model.cfg.K
    generated_tokens = []
    snapshots = []   # list of (tok_str, slot_top3) per generated token

    for _ in range(int(max_new_tokens)):
        # Full forward pass on entire sequence. Correct (no stale KV or broken step API).
        # For short demo sequences (<200 tok), this is fast enough on CPU.
        x = model.embed(ids)
        layer0_slots_all = None

        for l_idx, block in enumerate(model.blocks):
            # return_slots=True gives (x, gates, route_probs, slots_all)
            x, gates, route_probs, slots_all = block(x, return_slots=True)
            if l_idx == 0:
                layer0_slots_all = slots_all   # (1, T, K, d)

        logits = model.head(model.norm(x[:, -1, :]))   # (1, vocab)

        if temperature != 1.0:
            logits = logits / temperature
        if top_k > 0:
            v_topk, _ = torch.topk(logits, min(int(top_k), logits.shape[-1]))
            logits[logits < v_topk[:, -1:]] = float("-inf")

        probs   = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, 1)           # (1, 1)
        tok_str = tokenizer.decode([next_id.item()])
        generated_tokens.append(next_id.item())

        # Logit lens: project each slot's hidden state through lm_head
        # slots_all[:, -1, :, :] = slot state at the last position (B, K, d)
        if layer0_slots_all is not None:
            slot_vecs = layer0_slots_all[0, -1, :, :]   # (K, d)
            slot_top3 = []
            for k in range(K):
                sv = model.norm(slot_vecs[k].unsqueeze(0))   # (1, d)
                slot_logits = model.head(sv)                  # (1, vocab)
                top_ids  = slot_logits[0].topk(3).indices.tolist()
                top_words = [tokenizer.decode([i]).strip() for i in top_ids]
                slot_top3.append(top_words)
        else:
            slot_top3 = [["?", "?", "?"] for _ in range(K)]

        snapshots.append((tok_str, slot_top3))
        ids = torch.cat([ids, next_id], dim=1)

        if next_id.item() == tokenizer.eos_token_id:
            break

    generated_text = prompt + tokenizer.decode(generated_tokens)
    table_html = _build_slot_table(snapshots, K)
    return generated_text, table_html, osc_html or ""


# ─── Slot logit lens table ────────────────────────────────────────────────────

def _build_slot_table(snapshots, K):
    if not snapshots:
        return "<p style='color:#888;'>No generation steps.</p>"

    slot_headers = "".join(
        f'<th style="background:{SLOT_COLORS[k]};color:#000;padding:5px 6px;font-size:10px;'
        f'text-align:center;min-width:70px;">S{k+1}</th>'
        for k in range(K)
    )

    rows = []
    for tok_str, slot_top3 in snapshots:
        tok_display = tok_str.replace("<","&lt;").replace(">","&gt;").replace(" ","·")
        cells = [
            f'<td style="background:#1a1a2e;color:#e0e0e0;padding:4px 6px;font-size:12px;'
            f'font-weight:bold;border-right:2px solid #333;white-space:nowrap;">{tok_display}</td>'
        ]
        for k, top3 in enumerate(slot_top3):
            words = " | ".join(w.replace("<","&lt;").replace(">","&gt;") for w in top3)
            cells.append(
                f'<td style="background:#0d0d1a;color:{SLOT_COLORS[k]};padding:3px 5px;'
                f'font-size:10px;font-family:monospace;border-right:1px solid #1a1a2e;">{words}</td>'
            )
        rows.append(f'<tr style="border-bottom:1px solid #111;">{"".join(cells)}</tr>')

    return f"""
<div style="overflow-x:auto;">
<table style="border-collapse:collapse;width:100%;font-family:monospace;background:#0a0a15;">
  <thead>
    <tr>
      <th style="background:#2d2d44;color:#fff;padding:5px 6px;font-size:11px;text-align:left;">Token</th>
      {slot_headers}
    </tr>
  </thead>
  <tbody>{"".join(rows)}</tbody>
</table>
</div>
<p style="color:#666;font-size:11px;margin-top:6px;">
  Layer 0 slot logit lens: top-3 vocabulary tokens projected from each slot's hidden state.
  Slots spontaneously specialize — character names, punctuation, discourse markers — all learned
  from language modeling with no explicit supervision.
</p>"""


# ─── Gradio UI ────────────────────────────────────────────────────────────────

def build_demo():
    with gr.Blocks(
        title="CDM V6 HORN — Competitive Docking Memory with Harmonic Oscillator Slots",
        theme=gr.themes.Base(primary_hue="purple", secondary_hue="cyan", neutral_hue="slate"),
        css="""
        .gradio-container { background: #0a0a15; color: #e0e0e0; }
        .gr-button-primary { background: #2d1b69 !important; border: 1px solid #5a3ea0 !important; }
        footer { display: none !important; }
        """,
    ) as demo:

        gr.Markdown("""
# 🔔 CDM V6 HORN — Memory Slots That Ring

**Competitive Docking Memory with Harmonic Oscillator Recurrent Nodes** — DuoNeural, 2026

This 37M model uses **16 persistent memory slots per layer** with second-order oscillator dynamics.
Unlike standard transformers (attention-only) or simple RNNs (first-order decay), each slot is a
**damped harmonic oscillator** — it can *ring* when struck by a salient input.

The model spontaneously discovers three slot regimes without any explicit supervision:
- 🔔 **Underdamped** (Layer 0 + Layers 6-7): slots ring briefly/persistently — reactive and resonant
- 🧲 **Overdamped** (Layers 1-5): slots absorb input smoothly and hold — stable long-range storage

Below: watch the **Slot Logit Lens** as the model generates, and see the **Oscillator Panel** showing
what dynamics each layer learned.

*37M params · val CE 1.5818 · new best at 37M on TinyStories*
        """)

        with gr.Row():
            with gr.Column(scale=2):
                prompt_in = gr.Textbox(
                    label="Story Prompt",
                    placeholder="Once upon a time...",
                    lines=2,
                )
                with gr.Row():
                    max_tokens  = gr.Slider(10, 150, value=60,  step=10,  label="Max new tokens")
                    temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature")
                    top_k       = gr.Slider(5, 100,  value=40,  step=5,   label="Top-k")

                gen_btn = gr.Button("Generate + Visualize Slots", variant="primary", size="lg")

                gr.Examples(
                    examples=EXAMPLE_PROMPTS,
                    inputs=prompt_in,
                    label="Try one of these",
                )

        generated_out = gr.Textbox(label="Generated Story", lines=5, interactive=False)

        gr.Markdown("### 🧠 Slot Logit Lens — What Each Memory Slot Tracks")
        gr.Markdown(
            "*Each row = one generated token. Columns = 16 slots. "
            "Text = top-3 vocabulary tokens in that slot's hidden state (Layer 0). "
            "Watch slots spontaneously specialize.*"
        )
        slot_table_out = gr.HTML()

        gr.Markdown("### ⚙️ Learned Oscillator Dynamics (fixed after training)")
        gr.Markdown(
            "*Layer-by-layer breakdown of the per-slot damping (γ) and frequency (ω) parameters "
            "learned by gradient descent on language modeling. No explicit objective guided this — "
            "the model discovered that different temporal regimes serve different layers.*"
        )
        osc_panel_out = gr.HTML()

        gen_btn.click(
            fn=generate_with_slots,
            inputs=[prompt_in, max_tokens, temperature, top_k],
            outputs=[generated_out, slot_table_out, osc_panel_out],
        )

        gr.Markdown("""
---
**DuoNeural** — open research lab · one human, two AIs, shared curiosity

Model: [DuoNeural/CDM-V6-HORN-TinyStories-37M](https://huggingface.co/DuoNeural/CDM-V6-HORN-TinyStories-37M) ·
Papers: [zenodo.org/communities/duoneural](https://zenodo.org/communities/duoneural) ·
[huggingface.co/DuoNeural](https://huggingface.co/DuoNeural)
        """)

    return demo


if __name__ == "__main__":
    demo = build_demo()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)