app.py

"""
TemporalMesh Transformer (TMT) — Interactive Demo
Author: Vigneshwar LK · 2026 · MIT License
Paper: https://zenodo.org/records/20287390
GitHub: https://github.com/vignesh2027/TemporalMesh-Transformer
"""

import math
import random
import time
import warnings

import gradio as gr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots

warnings.filterwarnings("ignore")

# ─── Colour palette ────────────────────────────────────────────────────────────
BG   = "#0d1117"
SURF = "#161b22"
ACC  = "#58a6ff"
GRN  = "#3fb950"
PRP  = "#d2a8ff"
ORG  = "#ffa657"
RED  = "#ff7b72"
MUT  = "#8b949e"

# ══════════════════════════════════════════════════════════════════════════════
#  MINIMAL TMT IMPLEMENTATION — pure numpy, no torch needed for demo
# ══════════════════════════════════════════════════════════════════════════════

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-np.clip(x, -20, 20)))

def cosine_sim(a, b):
    na = np.linalg.norm(a) + 1e-8
    nb = np.linalg.norm(b) + 1e-8
    return float(np.dot(a, b) / (na * nb))

def build_mesh(X, k=4):
    """Build kNN graph from token representations. X: (S, D)"""
    S, D = X.shape
    k = min(k, S - 1)
    norms = np.linalg.norm(X, axis=1, keepdims=True) + 1e-8
    Xn = X / norms
    sim = Xn @ Xn.T                       # (S, S)
    np.fill_diagonal(sim, -2.0)           # exclude self
    edges = []
    weights = []
    for i in range(S):
        top_k = np.argsort(sim[i])[-k:]
        for j in top_k:
            edges.append((i, int(j)))
            weights.append(float(sim[i, j]))
    return edges, weights

def temporal_decay(i, j, S, w_decay=0.5):
    t_i, t_j = i / max(S - 1, 1), j / max(S - 1, 1)
    return float(sigmoid(w_decay * abs(t_i - t_j)))

def tmt_forward_demo(tokens, n_layers=6, k=4, exit_threshold=0.85, d=32, seed=42):
    """
    Minimal demo forward pass — returns per-layer diagnostics.
    tokens: list of strings
    """
    rng = np.random.RandomState(seed)
    S = len(tokens)
    # Random embeddings (fixed seed = deterministic)
    X = rng.randn(S, d).astype(np.float32)
    X = X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-8)

    exit_layer   = [n_layers] * S          # which layer each token exits
    confidences  = []                      # (n_layers, S)
    graphs       = []                      # (n_layers,) list of edge lists
    frozen       = np.zeros(S, dtype=bool)

    for layer in range(n_layers):
        # 1. Build mesh
        edges, wts = build_mesh(X, k=k)
        graphs.append((edges, wts))

        # 2. Sparse mesh attention (simplified)
        attn_out = X.copy()
        for i in range(S):
            if frozen[i]:
                continue
            nb_edges = [(j, w) for (src, j), w in zip(edges, wts) if src == i]
            if nb_edges:
                neighbours = [j for j, _ in nb_edges]
                raw        = np.array([wts[k2] * temporal_decay(i, j, S)
                                       for k2, (j, _) in enumerate(nb_edges)])
                alpha      = softmax(raw)
                agg        = sum(a * X[j] for a, (j, _) in zip(alpha, nb_edges))
                attn_out[i] = 0.7 * X[i] + 0.3 * agg

        # 3. Exit gate
        W_gate = rng.randn(d).astype(np.float32) * 0.3
        conf   = sigmoid(attn_out @ W_gate + 0.1 * layer)
        confidences.append(conf.tolist())

        for i in range(S):
            if not frozen[i] and conf[i] > exit_threshold:
                frozen[i] = True
                exit_layer[i] = layer + 1

        X = attn_out + rng.randn(S, d).astype(np.float32) * 0.02

    avg_depth    = float(np.mean(exit_layer))
    exit_pct     = float(np.sum(np.array(exit_layer) < n_layers) / S * 100)
    return {
        "exit_layer"    : exit_layer,
        "confidences"   : confidences,
        "graphs"        : graphs,
        "avg_depth"     : avg_depth,
        "exit_pct"      : exit_pct,
        "n_layers"      : n_layers,
        "tokens"        : tokens,
    }

# ══════════════════════════════════════════════════════════════════════════════
#  BENCHMARK DATA
# ══════════════════════════════════════════════════════════════════════════════

ABLATION = pd.DataFrame([
    {"Configuration": "Vanilla Transformer",  "Mesh": "✗", "Decay": "✗", "Exit": "✗",
     "Val PPL↓": 42.1, "Avg Layers": 12.0, "Rel Compute %": 100, "ΔPPL": "—"},
    {"Configuration": "Mesh Attention Only",  "Mesh": "✓", "Decay": "✗", "Exit": "✗",
     "Val PPL↓": 37.8, "Avg Layers": 12.0, "Rel Compute %":  62, "ΔPPL": "-10.2%"},
    {"Configuration": "Temporal Decay Only",  "Mesh": "✗", "Decay": "✓", "Exit": "✗",
     "Val PPL↓": 40.3, "Avg Layers": 12.0, "Rel Compute %":  98, "ΔPPL": "-4.3%"},
    {"Configuration": "Adaptive Exit Only",   "Mesh": "✗", "Decay": "✗", "Exit": "✓",
     "Val PPL↓": 39.6, "Avg Layers":  5.8, "Rel Compute %":  51, "ΔPPL": "-5.9%"},
    {"Configuration": "Mesh + Decay",         "Mesh": "✓", "Decay": "✓", "Exit": "✗",
     "Val PPL↓": 34.2, "Avg Layers": 12.0, "Rel Compute %":  61, "ΔPPL": "-18.8%"},
    {"Configuration": "Mesh + Exit",          "Mesh": "✓", "Decay": "✗", "Exit": "✓",
     "Val PPL↓": 35.1, "Avg Layers":  5.7, "Rel Compute %":  50, "ΔPPL": "-16.6%"},
    {"Configuration": "Decay + Exit",         "Mesh": "✗", "Decay": "✓", "Exit": "✓",
     "Val PPL↓": 37.0, "Avg Layers":  5.9, "Rel Compute %":  50, "ΔPPL": "-12.1%"},
    {"Configuration": "Full TMT (all 3)",     "Mesh": "✓", "Decay": "✓", "Exit": "✓",
     "Val PPL↓": 29.4, "Avg Layers":  5.5, "Rel Compute %":  48, "ΔPPL": "-30.2%"},
])

# ══════════════════════════════════════════════════════════════════════════════
#  PLOT HELPERS
# ══════════════════════════════════════════════════════════════════════════════

def _dark_fig(figsize=(10, 5)):
    fig, ax = plt.subplots(figsize=figsize, facecolor=BG)
    ax.set_facecolor(SURF)
    for spine in ax.spines.values():
        spine.set_edgecolor(MUT)
    ax.tick_params(colors=MUT)
    ax.xaxis.label.set_color(MUT)
    ax.yaxis.label.set_color(MUT)
    ax.title.set_color("#e6edf3")
    return fig, ax

def plot_ablation_bar():
    fig = go.Figure()
    cfgs  = ABLATION["Configuration"].tolist()
    ppls  = ABLATION["Val PPL↓"].tolist()
    comps = ABLATION["Rel Compute %"].tolist()
    colours = [GRN if "Full TMT" in c else ACC for c in cfgs]

    fig.add_trace(go.Bar(
        name="Validation PPL ↓", x=cfgs, y=ppls, marker_color=colours,
        text=[f"{v:.1f}" for v in ppls], textposition="outside",
        textfont=dict(color="#e6edf3", size=11),
    ))
    fig.add_trace(go.Scatter(
        name="Compute %", x=cfgs, y=comps, yaxis="y2",
        mode="lines+markers", line=dict(color=ORG, width=2.5, dash="dash"),
        marker=dict(size=8, color=ORG),
    ))
    fig.update_layout(
        paper_bgcolor=BG, plot_bgcolor=SURF,
        font=dict(color="#e6edf3", size=11),
        title=dict(text="Ablation Study — WikiText-2 (120M params)",
                   font=dict(size=16, color="#e6edf3"), x=0.5),
        yaxis=dict(title="Val Perplexity ↓", gridcolor="#30363d", color=MUT,
                   range=[25, 46]),
        yaxis2=dict(title="Relative Compute (%)", overlaying="y", side="right",
                    gridcolor="#30363d", color=ORG, range=[30, 115]),
        xaxis=dict(gridcolor="#30363d", color=MUT, tickangle=-20),
        legend=dict(bgcolor=SURF, bordercolor="#30363d"),
        height=420, margin=dict(l=60, r=60, t=60, b=80),
    )
    return fig

def plot_pareto():
    fig = go.Figure()
    ppls  = ABLATION["Val PPL↓"].tolist()
    comps = ABLATION["Rel Compute %"].tolist()
    cfgs  = ABLATION["Configuration"].tolist()
    colours = [GRN if "Full TMT" in c else ACC for c in cfgs]
    sizes   = [18 if "Full TMT" in c else 12 for c in cfgs]

    fig.add_trace(go.Scatter(
        x=comps, y=ppls, mode="markers+text",
        marker=dict(size=sizes, color=colours, line=dict(width=1, color="#30363d")),
        text=[c.replace(" (all 3)", "").replace(" Only", "") for c in cfgs],
        textposition="top center",
        textfont=dict(size=9, color="#e6edf3"),
    ))
    fig.add_annotation(
        x=48, y=29.4, text="← BEST (Full TMT)", showarrow=True,
        arrowhead=2, arrowcolor=GRN, font=dict(color=GRN, size=11),
        ax=-60, ay=-30,
    )
    fig.update_layout(
        paper_bgcolor=BG, plot_bgcolor=SURF,
        font=dict(color="#e6edf3", size=11),
        title=dict(text="Quality–Efficiency Pareto Frontier",
                   font=dict(size=16, color="#e6edf3"), x=0.5),
        xaxis=dict(title="Relative Compute (%) ↓", gridcolor="#30363d", color=MUT,
                   autorange="reversed"),
        yaxis=dict(title="Val PPL ↓", gridcolor="#30363d", color=MUT,
                   autorange="reversed"),
        height=400, margin=dict(l=60, r=40, t=60, b=60),
        showlegend=False,
    )
    return fig

def plot_complexity():
    S_vals = np.array([64, 128, 256, 512, 1024, 2048, 4096])
    k      = 8
    std_   = S_vals ** 2
    tmt_   = S_vals * k
    hier_  = S_vals ** (4/3)

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=S_vals, y=std_,  name="Standard O(S²)",
        mode="lines+markers", line=dict(color=RED, width=2.5),
        marker=dict(size=7, color=RED)))
    fig.add_trace(go.Scatter(x=S_vals, y=tmt_,  name=f"TMT Mesh O(S·k) k={k}",
        mode="lines+markers", line=dict(color=GRN, width=2.5),
        marker=dict(size=7, color=GRN)))
    fig.add_trace(go.Scatter(x=S_vals, y=hier_, name="Hierarchical O(S^4/3)",
        mode="lines+markers", line=dict(color=PRP, width=2.5, dash="dash"),
        marker=dict(size=7, color=PRP)))
    fig.update_layout(
        paper_bgcolor=BG, plot_bgcolor=SURF,
        font=dict(color="#e6edf3", size=11),
        title=dict(text="Attention Operations vs Sequence Length",
                   font=dict(size=16, color="#e6edf3"), x=0.5),
        xaxis=dict(title="Sequence Length (S)", type="log", gridcolor="#30363d", color=MUT),
        yaxis=dict(title="Attention Operations", type="log", gridcolor="#30363d", color=MUT),
        legend=dict(bgcolor=SURF, bordercolor="#30363d"),
        height=380, margin=dict(l=60, r=40, t=60, b=60),
    )
    return fig

def plot_exit_depth():
    token_types = ["Punctuation", "Determiners", "Common Nouns",
                   "Rare Nouns", "Technical", "OOV/Rare"]
    avg_depths  = [2.1, 3.4, 5.8, 9.3, 11.2, 11.7]
    colours     = [GRN, GRN, ACC, ORG, RED, RED]

    fig = go.Figure(go.Bar(
        x=token_types, y=avg_depths,
        marker_color=colours,
        text=[f"{d:.1f}" for d in avg_depths],
        textposition="outside",
        textfont=dict(color="#e6edf3", size=11),
    ))
    fig.add_hline(y=5.5, line=dict(color=PRP, dash="dash", width=2),
                  annotation_text="Avg 5.5 layers (Full TMT)",
                  annotation_font=dict(color=PRP))
    fig.update_layout(
        paper_bgcolor=BG, plot_bgcolor=SURF,
        font=dict(color="#e6edf3", size=11),
        title=dict(text="Exit Gate — Avg Processing Depth by Token Type",
                   font=dict(size=16, color="#e6edf3"), x=0.5),
        xaxis=dict(gridcolor="#30363d", color=MUT),
        yaxis=dict(title="Avg Layers Used", gridcolor="#30363d", color=MUT,
                   range=[0, 13]),
        height=380, margin=dict(l=60, r=40, t=60, b=60),
        showlegend=False,
    )
    return fig

def plot_graph_evolution(result):
    """Show graph edges for first 3 layers in a mini scatter."""
    tokens = result["tokens"]
    S      = len(tokens)
    graphs = result["graphs"]
    n_show = min(3, len(graphs))

    fig, axes = plt.subplots(1, n_show, figsize=(4 * n_show, 3.5), facecolor=BG)
    if n_show == 1:
        axes = [axes]

    # Fixed positions on a circle
    angles = np.linspace(0, 2 * np.pi, S, endpoint=False)
    xs = np.cos(angles)
    ys = np.sin(angles)

    layer_names = [f"Layer {i+1}" for i in range(n_show)]
    for ax, layer_idx, lname in zip(axes, range(n_show), layer_names):
        ax.set_facecolor(SURF)
        edges, wts = graphs[layer_idx]
        for (i, j), w in zip(edges, wts):
            alpha = max(0.15, min(0.9, (w + 1) / 2))
            ax.plot([xs[i], xs[j]], [ys[i], ys[j]],
                    color=ACC, alpha=alpha, lw=1.2)
        ax.scatter(xs, ys, s=160, c=GRN, zorder=5, edgecolors=BG, linewidths=1.5)
        for i, tok in enumerate(tokens):
            ax.text(xs[i] * 1.22, ys[i] * 1.22, tok[:6],
                    ha="center", va="center", fontsize=6.5,
                    color="#e6edf3")
        ax.set_title(lname, color="#e6edf3", fontsize=11, pad=6)
        ax.axis("off")
        for sp in ax.spines.values():
            sp.set_edgecolor(MUT)

    fig.suptitle("Dynamic Graph Topology Across Layers", color="#e6edf3",
                 fontsize=13, y=1.02)
    plt.tight_layout()
    return fig

def plot_exit_heatmap(result):
    tokens = result["tokens"]
    confs  = result["confidences"]  # (n_layers, S)
    n_layers = result["n_layers"]

    data = np.array(confs)  # (n_layers, S)
    fig, ax = plt.subplots(figsize=(max(6, len(tokens) * 0.7), 3.5), facecolor=BG)
    ax.set_facecolor(SURF)
    im = ax.imshow(data, cmap="YlOrRd", aspect="auto", vmin=0, vmax=1)
    plt.colorbar(im, ax=ax, label="Gate Confidence", fraction=0.03)
    im.colorbar.ax.yaxis.label.set_color(MUT)
    im.colorbar.ax.tick_params(colors=MUT)
    ax.set_xticks(range(len(tokens)))
    ax.set_xticklabels(tokens, rotation=35, ha="right", color=MUT, fontsize=8)
    ax.set_yticks(range(n_layers))
    ax.set_yticklabels([f"L{i+1}" for i in range(n_layers)], color=MUT, fontsize=8)
    ax.set_xlabel("Token", color=MUT)
    ax.set_ylabel("Layer", color=MUT)
    ax.set_title("Exit Gate Confidence per Token per Layer\n(darker = token froze here)",
                 color="#e6edf3", fontsize=11)

    # Mark exit points
    for s, el in enumerate(result["exit_layer"]):
        if el < n_layers:
            ax.add_patch(plt.Rectangle((s - 0.5, el - 0.5), 1, 1,
                         fill=False, edgecolor=GRN, lw=2))
    plt.tight_layout()
    return fig

# ══════════════════════════════════════════════════════════════════════════════
#  GRADIO CALLBACKS
# ══════════════════════════════════════════════════════════════════════════════

DEFAULT_TEXT = "The TemporalMesh Transformer achieves lower perplexity at reduced compute ."

def run_demo(text, n_layers, k, threshold, seed):
    if not text.strip():
        text = DEFAULT_TEXT
    tokens = text.strip().split()[:20]
    if len(tokens) < 3:
        tokens = DEFAULT_TEXT.split()[:10]

    result = tmt_forward_demo(
        tokens, n_layers=int(n_layers), k=int(k),
        exit_threshold=float(threshold), seed=int(seed),
    )

    fig_graph = plot_graph_evolution(result)
    fig_heat  = plot_exit_heatmap(result)

    exit_lyr  = result["exit_layer"]
    avg_depth = result["avg_depth"]
    exit_pct  = result["exit_pct"]

    summary_rows = []
    for i, (tok, el) in enumerate(zip(tokens, exit_lyr)):
        status = "Exited early" if el < n_layers else "Full depth"
        summary_rows.append({
            "Token": tok,
            "Exit Layer": el,
            "Status": status,
            "Compute Saved": f"{(1 - el/n_layers)*100:.0f}%",
        })
    df_summary = pd.DataFrame(summary_rows)

    stats = (
        f"**Tokens processed:** {len(tokens)}  |  "
        f"**Avg depth:** {avg_depth:.2f} / {n_layers} layers  |  "
        f"**Early exits:** {exit_pct:.0f}% of tokens  |  "
        f"**Compute vs full-depth:** {avg_depth/n_layers*100:.0f}%"
    )

    return fig_graph, fig_heat, df_summary, stats

def make_bench_tab():
    return (
        plot_ablation_bar(),
        plot_pareto(),
        plot_complexity(),
        plot_exit_depth(),
        ABLATION,
    )

# ══════════════════════════════════════════════════════════════════════════════
#  CSS
# ══════════════════════════════════════════════════════════════════════════════

CSS = """
body, .gradio-container { background: #0d1117 !important; color: #e6edf3 !important; }
.gr-button { background: #58a6ff !important; color: #0d1117 !important; font-weight: 700 !important; }
.gr-button:hover { background: #79c0ff !important; }
h1,h2,h3 { color: #e6edf3 !important; }
.gr-tab-nav button { color: #8b949e !important; }
.gr-tab-nav button.selected { color: #58a6ff !important; border-bottom: 2px solid #58a6ff !important; }
.gr-form { background: #161b22 !important; border: 1px solid #30363d !important; border-radius: 8px !important; }
.gr-input, .gr-textarea { background: #1c2333 !important; color: #e6edf3 !important;
  border: 1px solid #30363d !important; border-radius: 6px !important; }
.metric-banner { display: flex; gap: 1rem; flex-wrap: wrap; margin: 1rem 0; }
.metric-card { background: #161b22; border: 1px solid #30363d; border-radius: 10px;
  padding: 1rem 1.5rem; flex: 1; min-width: 140px; text-align: center; }
.metric-val { font-size: 2rem; font-weight: 900; }
.metric-lbl { font-size: .75rem; color: #8b949e; margin-top: .2rem; }
"""

HERO_MD = """
<div style="text-align:center;padding:2rem 1rem 1rem;">
  <div style="display:inline-flex;align-items:center;gap:.5rem;background:rgba(88,166,255,.1);
    border:1px solid rgba(88,166,255,.3);color:#58a6ff;font-size:.72rem;font-weight:700;
    letter-spacing:.1em;text-transform:uppercase;padding:.3rem 1rem;border-radius:999px;
    margin-bottom:1rem;">⬡ Novel Architecture · May 2026 · Open Source</div>
  <h1 style="font-size:2.8rem;font-weight:900;letter-spacing:-2px;margin:.5rem 0;
    background:linear-gradient(135deg,#e6edf3 0%,#58a6ff 45%,#d2a8ff 80%);
    -webkit-background-clip:text;-webkit-text-fill-color:transparent;">
    TemporalMesh Transformer</h1>
  <p style="color:#8b949e;font-size:1.05rem;margin:.5rem 0;">
    Dynamic Graph Attention · Temporal Decay · Adaptive Depth Routing</p>
  <div style="display:flex;gap:.75rem;justify-content:center;flex-wrap:wrap;margin:1rem 0;">
    <a href="https://zenodo.org/records/20287390" target="_blank"
      style="background:#58a6ff;color:#0d1117;padding:.45rem 1rem;border-radius:7px;
      text-decoration:none;font-weight:700;font-size:.88rem;">📄 Paper</a>
    <a href="https://github.com/vignesh2027/TemporalMesh-Transformer" target="_blank"
      style="background:#30363d;color:#e6edf3;padding:.45rem 1rem;border-radius:7px;
      text-decoration:none;font-weight:700;font-size:.88rem;">⭐ GitHub</a>
    <a href="https://huggingface.co/vigneshwar234/TemporalMesh-Transformer" target="_blank"
      style="background:#30363d;color:#e6edf3;padding:.45rem 1rem;border-radius:7px;
      text-decoration:none;font-weight:700;font-size:.88rem;">🤗 Model</a>
  </div>
  <div style="display:grid;grid-template-columns:repeat(auto-fit,minmax(130px,1fr));
    gap:1px;background:#30363d;border:1px solid #30363d;border-radius:10px;
    overflow:hidden;max-width:700px;margin:1.5rem auto 0;">
    <div style="background:#161b22;padding:1rem;text-align:center;">
      <div style="font-size:1.9rem;font-weight:900;color:#3fb950;">29.4</div>
      <div style="font-size:.72rem;color:#8b949e;">TMT PPL (WikiText-2)</div>
    </div>
    <div style="background:#161b22;padding:1rem;text-align:center;">
      <div style="font-size:1.9rem;font-weight:900;color:#ffa657;">42.1</div>
      <div style="font-size:.72rem;color:#8b949e;">Vanilla Baseline PPL</div>
    </div>
    <div style="background:#161b22;padding:1rem;text-align:center;">
      <div style="font-size:1.9rem;font-weight:900;color:#58a6ff;">30.2%</div>
      <div style="font-size:.72rem;color:#8b949e;">PPL Reduction</div>
    </div>
    <div style="background:#161b22;padding:1rem;text-align:center;">
      <div style="font-size:1.9rem;font-weight:900;color:#d2a8ff;">0.48×</div>
      <div style="font-size:.72rem;color:#8b949e;">Relative Compute</div>
    </div>
    <div style="background:#161b22;padding:1rem;text-align:center;">
      <div style="font-size:1.9rem;font-weight:900;color:#3fb950;">226</div>
      <div style="font-size:.72rem;color:#8b949e;">Tests Passing</div>
    </div>
  </div>
</div>
"""

ARCH_MD = """
## How TMT Works

TMT introduces **three tightly coupled innovations** on top of a standard transformer:

### ⬡ Innovation 1 — Mesh Attention
Instead of every token attending to every other token (O(S²)), TMT builds a **dynamic kNN graph** from cosine similarity of current token representations. Only the top-k nearest neighbours exchange attention. The graph **rebuilds every layer** — topology is emergent, not fixed.

```
edge_index = topk(cosine_sim(X), k=8)     # per token, per layer
attn = softmax(QKᵀ/√d) × cosine_weight    # sparse: O(S·k) vs O(S²)
```

### ⏱ Innovation 2 — Temporal Decay
A learned per-head scalar multiplies into **post-softmax attention weights**, attenuating semantically distant tokens. Unlike ALiBi (additive to logits, fixed), TMT's decay is:
- **Multiplicative** (applied after softmax)
- **Learned** (per head, not fixed schedule)
- **Semantic** (models relevance, not just position)

```
δ_h(i,j) = sigmoid(W_decay_h × |t_i − t_j|)    # W_decay learned
ã_ij = α_ij × δ_h(i,j)                          # post-softmax multiply
```

### ⚡ Innovation 3 — Adaptive Depth Routing
After each layer, a gate evaluates each token's confidence. If `confidence > τ=0.85`, the token **freezes** and skips all remaining layers — saving compute.

```
conf = sigmoid(W_gate · x)    # per-token scalar
if conf > 0.85: token.freeze()    # exits here
else: continue to next layer
```

**Result:** Average 5.5 layers used vs 12. 2.1× compute efficiency with *lower* perplexity.

### Plus: Dual-Stream FFN + EMA Memory Anchors
- **DualStream FFN:** Two parallel 256-dim streams (syntax + semantic) fused by a learned gate
- **Memory Anchors:** 16 persistent EMA-updated parameter vectors providing global cross-sequence context

---

### Forward Pass Summary
```
input_ids (B, S)
    → TokenEmbedding (B, S, D)
    → TemporalPositionEncoder (RoPE + decay scalars)
    → MeshBuilder → kNN graph (edge_index, weights)
    → N=12 × TMTLayer:
        ├── MeshAttention (sparse, graph-constrained)
        ├── DualStreamFFN (syntax + semantic)
        ├── ExitGate (per-token confidence)
        └── MemoryAnchorCross (16 EMA anchors)
    → LayerNorm → OutputProjection
    → TMTOutput { logits, exit_masks, confidences, graph_edges, memory_state, decay_scalars }
```
"""

CITE_MD = """
## Citation

```bibtex
@article{vigneshwar2026temporalmesh,
  title   = {TemporalMesh Transformer: Dynamic Graph Attention
             with Temporal Decay and Adaptive Depth Routing},
  author  = {LK, Vigneshwar},
  journal = {Zenodo Preprint},
  year    = {2026},
  doi     = {10.5281/zenodo.20287197},
  url     = {https://zenodo.org/records/20287390}
}
```

## Quick Install
```bash
git clone https://github.com/vignesh2027/TemporalMesh-Transformer.git
cd TemporalMesh-Transformer
pip install torch einops && pip install -e .
```

```python
from tmt.model.config import TMTConfig
from tmt.model.model  import TMTModel
import torch

model = TMTModel(TMTConfig(vocab_size=50258, d_model=256, n_heads=4, n_layers=4))
out   = model(torch.randint(0, 50258, (1, 64)))

print(out.logits.shape)       # (1, 64, 50258)
print(out.memory_state.shape) # (16, 256)
print(out.exit_masks[-1].float().mean())  # fraction exited early
```

## Links
- 📄 **Paper:** [zenodo.org/records/20287390](https://zenodo.org/records/20287390)
- ⭐ **GitHub:** [github.com/vignesh2027/TemporalMesh-Transformer](https://github.com/vignesh2027/TemporalMesh-Transformer)
- 🤗 **Model:** [huggingface.co/vigneshwar234/TemporalMesh-Transformer](https://huggingface.co/vigneshwar234/TemporalMesh-Transformer)
- 📊 **Dataset:** [huggingface.co/datasets/vigneshwar234/TMT-Benchmarks](https://huggingface.co/datasets/vigneshwar234/TMT-Benchmarks)
- **DOI:** 10.5281/zenodo.20287197

**Author:** Vigneshwar LK · Independent Research · 2026 · MIT License
"""

# ══════════════════════════════════════════════════════════════════════════════
#  BUILD APP
# ══════════════════════════════════════════════════════════════════════════════

with gr.Blocks(css=CSS, title="TemporalMesh Transformer Demo") as demo:
    gr.HTML(HERO_MD)

    with gr.Tabs():

        # ── Tab 1: Live Demo ─────────────────────────────────────────────────
        with gr.Tab("⚡ Live Forward Pass"):
            gr.Markdown("### Run a TMT forward pass — watch the dynamic graph and exit gates live")
            with gr.Row():
                with gr.Column(scale=2):
                    txt_input  = gr.Textbox(
                        label="Input text (space-separated tokens, max 20)",
                        value=DEFAULT_TEXT,
                        lines=2, placeholder="Enter your sentence here…"
                    )
                    with gr.Row():
                        sl_layers = gr.Slider(2, 12, value=6, step=1,   label="Layers (n_layers)")
                        sl_k      = gr.Slider(2,  8, value=4, step=1,   label="Graph k (neighbours)")
                    with gr.Row():
                        sl_thresh = gr.Slider(0.5, 0.99, value=0.85, step=0.01, label="Exit threshold τ")
                        sl_seed   = gr.Slider(0, 100, value=42, step=1,          label="Random seed")
                    btn_run = gr.Button("▶  Run TMT Forward Pass", variant="primary")
                    md_stats = gr.Markdown("Click **Run** to start…")
                with gr.Column(scale=1):
                    gr.Markdown("""
**What you'll see:**

🔵 **Graph panels** — dynamic kNN topology across first 3 layers. Edges connect semantically similar tokens; density and specialisation change with depth.

🔴 **Heatmap** — exit gate confidence per token per layer. Green box = token froze at that layer.

📊 **Table** — per-token exit layer and compute saved.

> With τ=0.85 and 6 layers, ~50% of tokens typically exit before the final layer — especially punctuation and common words.
""")

            with gr.Row():
                plot_graph_out = gr.Plot(label="Dynamic Graph Topology")
                plot_heat_out  = gr.Plot(label="Exit Gate Confidence Heatmap")

            df_token_out = gr.Dataframe(
                label="Per-Token Exit Analysis",
                headers=["Token", "Exit Layer", "Status", "Compute Saved"],
            )

            btn_run.click(
                fn=run_demo,
                inputs=[txt_input, sl_layers, sl_k, sl_thresh, sl_seed],
                outputs=[plot_graph_out, plot_heat_out, df_token_out, md_stats],
            )

        # ── Tab 2: Benchmarks ────────────────────────────────────────────────
        with gr.Tab("📊 Benchmarks"):
            gr.Markdown("### Ablation Study — WikiText-2 · 120M parameters · 10,000 training steps")
            gr.Markdown("""
> **Key result:** Full TMT achieves **PPL 29.4 vs Vanilla 42.1** — a **30.2% reduction** — while using only **48% of compute**.
> The combined gain (12.7 PPL) exceeds the sum of individual components (8.6 PPL) — confirming **superadditive synergy**.
""")
            with gr.Row():
                plot_abl  = gr.Plot(label="Ablation Bar Chart")
                plot_par  = gr.Plot(label="Pareto Frontier")
            with gr.Row():
                plot_cplx = gr.Plot(label="Complexity Scaling")
                plot_exit = gr.Plot(label="Exit Depth by Token Type")

            df_abl = gr.Dataframe(label="Full Ablation Table", interactive=False)

            demo.load(
                fn=make_bench_tab,
                outputs=[plot_abl, plot_par, plot_cplx, plot_exit, df_abl],
            )

        # ── Tab 3: Architecture ──────────────────────────────────────────────
        with gr.Tab("🏗 Architecture"):
            gr.Markdown(ARCH_MD)

        # ── Tab 4: Citation / Links ──────────────────────────────────────────
        with gr.Tab("📄 Paper & Citation"):
            gr.Markdown(CITE_MD)

if __name__ == "__main__":
    demo.launch()