Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| app.py — HuggingFace Spaces demo for Competitive Docking Memory (CDM) Language Model | |
| DuoNeural — Archon + Jesse Caldwell + Aura — 2026 | |
| Gradio 5 compatible. Routing gate visualization (not Logit Lens). | |
| """ | |
| import json | |
| import torch | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| model = None | |
| tokenizer = None | |
| DEVICE = "cpu" | |
| MODEL_REPO = "DuoNeural/CDM-V2-TinyStories-37M" | |
| SLOT_COLORS = [ | |
| "#FF6B6B","#4ECDC4","#45B7D1","#96CEB4","#FECA57","#FF9FF3","#54A0FF","#5F27CD", | |
| "#00D2D3","#FF9F43","#C8D6E5","#8395A7","#EE5A24","#009432","#C4E538","#A3CB38", | |
| ] | |
| EXAMPLES = [ | |
| "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", | |
| "She wanted to play outside but the weather was", | |
| ] | |
| SLOT_ROLES = { | |
| 11: "PUNCT", # 0-indexed — slot 11 specializes on punctuation (step5000: 71%, step30k: 88%) | |
| 9: "AGENCY", # character agency verbs | |
| 0: "IDENTITY",# character names / pronouns | |
| 15: "VERBS", # action verbs | |
| 5: "ARTICLES",# articles & modifiers | |
| } | |
| def load_model(): | |
| global model, tokenizer | |
| if model is not None: | |
| return | |
| from transformers import GPT2TokenizerFast | |
| from cdm_model_v2 import CDMConfigV2, CDMLanguageModelV2 | |
| tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") | |
| tokenizer.pad_token = tokenizer.eos_token | |
| model_pt_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.pt") | |
| config_path = hf_hub_download(repo_id=MODEL_REPO, filename="config.json") | |
| with open(config_path) as f: | |
| cfg_dict = json.load(f) | |
| cfg = CDMConfigV2( | |
| 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), | |
| ) | |
| ckpt = torch.load(model_pt_path, map_location=DEVICE, weights_only=False) | |
| model = CDMLanguageModelV2(cfg) | |
| model.load_state_dict(ckpt.get("model_state", ckpt)) | |
| model = model.to(DEVICE) | |
| model.eval() | |
| print("Loading CDM model on startup...") | |
| try: | |
| load_model() | |
| print("Model loaded OK") | |
| except Exception as e: | |
| print(f"Startup load failed (will retry on first request): {e}") | |
| def gate_to_css_color(base_hex: str, intensity: float) -> str: | |
| """Blend base color with dark background by intensity (0-1).""" | |
| r = int(base_hex[1:3], 16) | |
| g = int(base_hex[3:5], 16) | |
| b = int(base_hex[5:7], 16) | |
| bg = 10 | |
| r2 = int(bg + (r - bg) * intensity) | |
| g2 = int(bg + (g - bg) * intensity) | |
| b2 = int(bg + (b - bg) * intensity) | |
| return f"#{r2:02x}{g2:02x}{b2:02x}" | |
| def generate_and_visualize(prompt: str, max_new_tokens: int, temperature: float, top_k: int): | |
| if model is None: | |
| try: | |
| load_model() | |
| except Exception as e: | |
| return f"[Model load failed: {e}]", "<p>Model not loaded.</p>", "" | |
| ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE) | |
| generated_text, snapshots = model.generate_with_slots( | |
| ids, max_new=int(max_new_tokens), tokenizer=tokenizer, | |
| temperature=float(temperature), top_k=int(top_k), | |
| ) | |
| K = model.cfg.K | |
| # ── Routing gate heatmap table ──────────────────────────────────────────── | |
| rows = [] | |
| slot_token_counts = [[] for _ in range(K)] # track which tokens route to each slot | |
| for tok_str, all_layer_gates, winner in snapshots: | |
| tok_display = tok_str.replace("<","<").replace(">",">").replace(" ","·").strip() or "⏎" | |
| last_gates = all_layer_gates[-1] # final layer gate distribution | |
| # Track routing for summary | |
| slot_token_counts[winner].append(tok_str.strip()) | |
| # Winner badge | |
| win_color = SLOT_COLORS[winner % len(SLOT_COLORS)] | |
| winner_label = SLOT_ROLES.get(winner, f"S{winner+1}") | |
| tok_cell = ( | |
| f'<td style="background:#1a1a2e;color:#e0e0e0;padding:3px 8px;' | |
| f'font-size:13px;font-weight:bold;border-right:2px solid #333;white-space:nowrap;">' | |
| f'{tok_display} ' | |
| f'<span style="color:{win_color};font-size:10px;font-weight:normal;">→{winner_label}</span>' | |
| f'</td>' | |
| ) | |
| gate_cells = [] | |
| for k, gv in enumerate(last_gates): | |
| intensity = min(1.0, gv * K) # scale up since each slot avg ~1/K | |
| bg = gate_to_css_color(SLOT_COLORS[k % len(SLOT_COLORS)], intensity) | |
| pct = f"{gv*100:.0f}%" if gv > 0.08 else "" | |
| border = f"2px solid {SLOT_COLORS[k%len(SLOT_COLORS)]}" if k == winner else "1px solid #111" | |
| gate_cells.append( | |
| f'<td style="background:{bg};color:#fff;padding:2px 5px;' | |
| f'font-size:10px;font-family:monospace;border:{border};' | |
| f'text-align:center;min-width:28px;">{pct}</td>' | |
| ) | |
| rows.append(f'<tr style="border-bottom:1px solid #111;">{tok_cell}{"".join(gate_cells)}</tr>') | |
| # Header | |
| slot_headers = "".join( | |
| f'<th style="background:{SLOT_COLORS[k%len(SLOT_COLORS)]};color:#000;' | |
| f'padding:4px 3px;font-size:10px;font-weight:bold;text-align:center;min-width:28px;">' | |
| f'{"★" if k in SLOT_ROLES else f"S{k+1}"}</th>' | |
| for k in range(K) | |
| ) | |
| table_html = f"""<div style="overflow-x:auto;background:#050510;border-radius:8px;padding:8px;"> | |
| <p style="color:#888;font-size:11px;margin:0 0 6px 0;"> | |
| Cell brightness = routing affinity. ★ = slots with known specialization (probe-confirmed). | |
| → label = winning slot for each token. <b>Watch S12 light up for punctuation!</b> (slot 11, 0-indexed) | |
| </p> | |
| <table style="border-collapse:collapse;width:100%;font-family:monospace;"> | |
| <thead><tr> | |
| <th style="background:#1a1a2e;color:#aaa;padding:4px 7px;font-size:11px;text-align:left;border-right:2px solid #333;white-space:nowrap;">Token → Winner</th> | |
| {slot_headers}</tr></thead> | |
| <tbody>{"".join(rows)}</tbody> | |
| </table></div>""" | |
| # ── Routing summary ─────────────────────────────────────────────────────── | |
| summary_lines = ["Slot routing summary (final layer, generated tokens):", ""] | |
| for k in range(K): | |
| tokens = slot_token_counts[k] | |
| role = SLOT_ROLES.get(k, "") | |
| role_str = f" [{role}]" if role else "" | |
| if tokens: | |
| top_toks = ", ".join(f'"{t}"' for t in tokens[:6] if t) | |
| summary_lines.append(f"S{k+1:2d}{role_str}: {len(tokens)} wins — {top_toks}") | |
| else: | |
| summary_lines.append(f"S{k+1:2d}{role_str}: 0 wins") | |
| summary = "\n".join(summary_lines) | |
| return generated_text, table_html, summary | |
| with gr.Blocks(title="CDM V2 — Competitive Docking Memory | DuoNeural") as demo: | |
| gr.Markdown(""" | |
| # 🧠 Competitive Docking Memory (CDM) — Routing Gate Visualization | |
| **[DuoNeural](https://huggingface.co/DuoNeural) 2026 | Archon · Jesse Caldwell · Aura** | |
| This 37M-parameter model maintains **K=16 persistent memory slots** per layer that compete to process each token. | |
| Slots **spontaneously specialize** without supervision: slot 11 (S12) dominates punctuation, slot 0 (S1) tracks character names, etc. | |
| The heatmap shows **routing gate affinity** — how strongly each slot claims each token. | |
| Bright cell = that slot "owns" the token. The `→ label` shows the winning slot. | |
| > *K_eff ≈ 15.9/16 · 99.8% Shannon Capacity Saturation · val CE 1.5934 vs 1.6516 vanilla GPT baseline* | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| prompt_in = gr.Textbox(label="Story Prompt", placeholder="Once upon a time...", | |
| lines=2, value="Once upon a time there was a little girl named Lily.") | |
| with gr.Row(): | |
| max_tokens = gr.Slider(10, 80, value=50, step=5, label="Max new tokens") | |
| temperature = gr.Slider(0.1, 1.4, value=0.85, step=0.05, label="Temperature") | |
| top_k = gr.Slider(5, 80, value=40, step=5, label="Top-k") | |
| btn = gr.Button("⚡ Generate + Show Routing Gates", variant="primary") | |
| with gr.Column(scale=2): | |
| routing_summary = gr.Textbox(label="Routing Summary — which tokens each slot claimed", | |
| lines=12, interactive=False) | |
| generated_out = gr.Textbox(label="Generated Story", lines=4, interactive=False) | |
| gr.Markdown("""### Routing Gate Heatmap | |
| *Each row = one generated token. Columns S1–S16 = memory slots (★ = probe-confirmed specializations). | |
| Cell brightness = gate affinity. `→ label` = winning slot. Try a prompt with punctuation to watch S12!*""") | |
| slot_table = gr.HTML() | |
| gr.Examples(examples=EXAMPLES, inputs=prompt_in, label="Try these prompts") | |
| btn.click(fn=generate_and_visualize, | |
| inputs=[prompt_in, max_tokens, temperature, top_k], | |
| outputs=[generated_out, slot_table, routing_summary]) | |
| gr.Markdown("""--- | |
| **Known slot specializations** (from routing probe, step 30000, last layer): | |
| S12 (slot 11) = PUNCT · S10 (slot 9) = CHARACTER AGENCY · S1 (slot 0) = CHARACTER IDENTITY · | |
| S16 (slot 15) = ACTION VERBS · S6 (slot 5) = ARTICLES | |
| **Architecture:** 37.1M params · d=384 · 8 layers · K=16 competitive memory slots · GQA · SwiGLU FFN | |
| **Model:** [DuoNeural/CDM-V2-TinyStories-37M](https://huggingface.co/DuoNeural/CDM-V2-TinyStories-37M) | |
| **Paper:** *Competitive Docking Memory: Emergent Slot Specialization in Language Models* — DuoNeural 2026""") | |
| demo.launch() | |