Update app.py

app.py

@@ -1,521 +1,271 @@
 from __future__ import annotations
 
+import os
+os.environ["KERAS_BACKEND"] = "jax"
+
 import numpy as np
-import keras
-from keras import layers, ops
 import jax
-
-from veylon_attention import flash_splash_attention, decode_swa
-
-try:
-    from config import (
-        CONTEXT,
-        vocab_size as Vocab_size,
-        D_MODEL,
-        numberoflayers,
-        numberofheads,
-        d_Latent,
-        ffn_mult,
-        swa_window,
-        num_kv_heads,
-        use_moe,
-        moe_num_experts,
-        moe_top_k,
-        USE_REMAT,
-        MAX_GEN_TOKENS,
-    )
-except Exception:
-    CONTEXT = 2048
-    Vocab_size = 32000
-    D_MODEL = 512
-    numberoflayers = 8
-    numberofheads = 8
-    d_Latent = 128
-    ffn_mult = 3.5
-    swa_window = 1024
-    num_kv_heads = 2
-    use_moe = False
-    moe_num_experts = 8
-    moe_top_k = 2
-    USE_REMAT = True
-    MAX_GEN_TOKENS = 256
-
-# NOTE: mixed precision policy is set by train.py / finetune.py before model
-# creation.  Do NOT set it here — importing this module must not have side
-# effects.
-
-
-# ─────────────────────────────────────────────────────────────────────────────
-# Layers
-# ─────────────────────────────────────────────────────────────────────────────
-
-@keras.saving.register_keras_serializable()
-class RMSNorm(layers.Layer):
-    def __init__(self, epsilon=1e-5, **kwargs):
-        super().__init__(**kwargs)
-        self.epsilon = epsilon
-
-    def build(self, input_shape):
-        self.weight = self.add_weight(shape=(input_shape[-1],), initializer='ones', name='gamma')
-
-    def call(self, x):
-        x_fp32 = ops.cast(x, 'float32')
-        rms = ops.sqrt(ops.mean(ops.square(x_fp32), axis=-1, keepdims=True) + self.epsilon)
-        out = x_fp32 / rms
-        return ops.cast(out, x.dtype) * self.weight
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({'epsilon': self.epsilon})
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class RotaryEmbedding(layers.Layer):
+import keras
+import gradio as gr
+from pathlib import Path
+
+from veylon_model import create_llm
+from tokenizer import TokenizerWrapper
+from config import (
+    CONTEXT,
+    vocab_size,
+    D_MODEL,
+    numberoflayers,
+    numberofheads,
+    d_Latent,
+    ffn_mult,
+    num_kv_heads,
+    swa_window,
+)
+
+# ============================================================
+# Initialize (runs once)
+# ============================================================
+
+keras.mixed_precision.set_global_policy("mixed_bfloat16")
+
+print(f"Backend: {keras.backend.backend()}")
+print(f"JAX devices: {jax.devices()}")
+
+# Load tokenizer
+tokenizer = TokenizerWrapper("tokenizer.model")
+assert tokenizer.vocab_size == vocab_size, (
+    f"Tokenizer vocab ({tokenizer.vocab_size}) != config vocab ({vocab_size})"
+)
+print(f"✓ Tokenizer loaded: {tokenizer.vocab_size} vocab")
+
+# Build model
+print("Building model...")
+model = create_llm(
+    vocab_size=vocab_size,
+    d_model=D_MODEL,
+    n_layers=numberoflayers,
+    n_heads=numberofheads,
+    d_latent=d_Latent,
+    ffn_mult=ffn_mult,
+    max_seq_len=CONTEXT,
+    use_moe=False,
+    num_kv_heads=num_kv_heads,
+    swa_window=swa_window,
+)
+
+# Warmup
+dummy = np.zeros((1, CONTEXT), dtype=np.int32)
+_ = model(dummy, training=False)
+print("✓ Model built successfully")
+
+# Load weights
+WEIGHTS_PATH = "veylon_final.weights.h5"
+if Path(WEIGHTS_PATH).exists():
+    print(f"Loading weights from: {WEIGHTS_PATH}")
+    model.load_weights(WEIGHTS_PATH)
+    print("✓ Weights loaded successfully")
+else:
+    print(f"WARNING: {WEIGHTS_PATH} not found. Using untrained model.")
+
+print(f"✓ Model params: {model.count_params():,}\n")
+
+# ============================================================
+# Sampling
+# ============================================================
+
+def sample_from_logits(
+    logits: np.ndarray,
+    temperature: float = 0.8,
+    top_k: int = 50,
+) -> int:
+    """NumPy-only sampling."""
+    logits = np.array(logits, dtype=np.float32, copy=True)
+
+    if temperature > 0:
+        logits = logits / float(max(temperature, 1e-8))
+
+    if top_k > 0:
+        k = min(int(top_k), logits.shape[-1])
+        row = logits[0]
+        top_indices = np.argpartition(row, -k)[-k:]
+        filtered = np.full_like(row, -np.inf)
+        filtered[top_indices] = row[top_indices]
+        logits[0] = filtered
+
+    row = logits[0]
+    row = row - np.max(row)
+    probs = np.exp(row)
+    probs = probs / probs.sum()
+
+    return int(np.random.choice(len(probs), p=probs))
+
+# ============================================================
+# Generation function
+# ============================================================
+
+def generate(
+    prompt: str,
+    max_new_tokens: int = 64,
+    temperature: float = 0.8,
+    top_k: int = 50,
+) -> str:
     """
-    RoPE with `max_seq_len + gen_headroom` so a full-context prompt can still
-    decode up to `gen_headroom` tokens without overflowing the cos/sin table.
-
-    `max_seq_len` is preserved in the config so checkpoints remain compatible —
-    only the (trainable=False) table shapes grow.
+    Generate text from a prompt using Veylon.
+    
+    Args:
+        prompt: Input text
+        max_new_tokens: Maximum tokens to generate
+        temperature: Sampling temperature (0.1-2.0)
+        top_k: Top-K sampling cutoff
+    
+    Returns:
+        Generated text
     """
+    try:
+        # Encode prompt
+        tokens = tokenizer.encode(
+            prompt,
+            add_bos=True,
+            add_eos=False,
+        )
+
+        if len(tokens) == 0:
+            tokens = [tokenizer.bos_id if hasattr(tokenizer, "bos_id") else 1]
+
+        tokens = tokens[-CONTEXT:]
 
-    def __init__(self, max_seq_len, dim, theta=10000.0, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        if dim % 2 != 0:
-            raise ValueError('RotaryEmbedding dim must be even.')
-        self.max_seq_len = max_seq_len
-        self.dim = dim
-        self.theta = theta
-        self.gen_headroom = int(gen_headroom)
-        self.table_size = max_seq_len + self.gen_headroom
-
-    def build(self, input_shape):
-        half = self.dim // 2
-        inv_freq = 1.0 / (self.theta ** (np.arange(0, self.dim, 2).astype(np.float32) / self.dim))
-        positions = np.arange(self.table_size, dtype=np.float32)
-        freqs = positions[:, None] * inv_freq[None, :]
-        self.cos = self.add_weight(
-            shape=(self.table_size, half),
-            initializer=keras.initializers.Constant(np.cos(freqs)),
-            trainable=False, dtype='float32', name='cos_table',
+        # Prefill phase
+        prompt_ids = np.array([tokens], dtype=np.int32)
+        logits, cache_k, cache_v = model.generate_step(
+            prompt_ids,
+            cache_k=None,
+            cache_v=None,
+            cache_pos=0,
         )
-        self.sin = self.add_weight(
-            shape=(self.table_size, half),
-            initializer=keras.initializers.Constant(np.sin(freqs)),
-            trainable=False, dtype='float32', name='sin_table',
+
+        next_token = sample_from_logits(
+            np.array(logits[:, -1, :], dtype=np.float32, copy=True),
+            temperature=temperature,
+            top_k=top_k,
         )
+        tokens.append(next_token)
+
+        # Decoding phase (token-by-token)
+        if next_token != tokenizer.eos_id and len(tokens) < CONTEXT:
+            cache_pos = len(prompt_ids[0])
 
-    def call(self, x, offset=0):
-        seq_len = x.shape[1]
-        if seq_len is None:
-            seq_len = ops.shape(x)[1]
-        half = self.dim // 2
-
-        # JIT FIX: `offset` may be a traced jnp array (e.g. cache_pos passed
-        # in under jax.jit for compiled autoregressive decoding), not a
-        # Python int. The original code did `if offset + seq_len >
-        # self.table_size: raise ...` and `self.cos[offset:offset+seq_len]`
-        # — both are illegal under tracing: you cannot branch a Python `if`
-        # on a traced value, and Python slice syntax requires a static start
-        # index. We replace the bounds check with a debug-only assertion
-        # that only fires when offset is a concrete Python int/np scalar
-        # (i.e. eager calls, not jitted ones — jitted calls trust the caller
-        # to pass a valid offset, same as the KV-cache contract elsewhere in
-        # this file), and replace the slice with ops.slice /
-        # dynamic_slice_in_dim, which correctly handles a traced start index.
-        if isinstance(offset, int):
-            if offset + seq_len > self.table_size:
-                raise ValueError(
-                    f'RoPE table too small: offset={offset}, seq_len={seq_len}, '
-                    f'max={self.table_size} (CONTEXT={self.max_seq_len} + gen_headroom={self.gen_headroom})'
+            for _ in range(max_new_tokens - 1):
+                next_input = np.array([[next_token]], dtype=np.int32)
+
+                logits, cache_k, cache_v = model.generate_step(
+                    next_input,
+                    cache_k=cache_k,
+                    cache_v=cache_v,
+                    cache_pos=cache_pos,
                 )
 
-        offset = ops.cast(offset, 'int32')
-        cos = ops.slice(self.cos, [offset, 0], [seq_len, half])
-        sin = ops.slice(self.sin, [offset, 0], [seq_len, half])
-        cos = ops.cast(ops.reshape(cos, (1, seq_len, 1, half)), x.dtype)
-        sin = ops.cast(ops.reshape(sin, (1, seq_len, 1, half)), x.dtype)
-        x1 = x[..., :half]
-        x2 = x[..., half:]
-        return ops.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'max_seq_len': self.max_seq_len,
-            'dim': self.dim,
-            'theta': self.theta,
-            'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class SwiGLUFFN(layers.Layer):
-    def __init__(self, d_model, hidden_mult=3.5, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model_arg = d_model
-        self.hidden_mult = hidden_mult
-        self.hidden_dim = int(d_model * hidden_mult * 2 / 3)
-        self.hidden_dim = ((self.hidden_dim + 63) // 64) * 64
-
-    def build(self, input_shape):
-        d_model = input_shape[-1]
-        self.gate_up_proj = self.add_weight(shape=(d_model, 2 * self.hidden_dim), initializer='glorot_uniform', name='gate_up_proj')
-        self.down_proj = self.add_weight(shape=(self.hidden_dim, d_model), initializer='glorot_uniform', name='down_proj')
-
-    def call(self, x, training=False):
-        gate_up = ops.matmul(x, self.gate_up_proj)
-        gate, up = ops.split(gate_up, 2, axis=-1)
-        return ops.matmul(ops.silu(gate) * up, self.down_proj)
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({'d_model': self.d_model_arg, 'hidden_mult': self.hidden_mult})
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class MoE_FFN(layers.Layer):
-    """
-    Mixture-of-Experts FFN (top-k routing, masked dispatch).
+                cache_pos += 1
 
-    For each routing slot k in [0, top_k) we run each expert once over the
-    full batch, zeroing inputs the expert was not selected for, then accumulate
-    the gated expert output.  The load-balancing loss uses top-1 argmax routing
-    for the `f` term (standard Switch-Transformer formulation); the previous
-    code used top-k membership which over-counted and produced the wrong loss.
-    """
+                next_token = sample_from_logits(
+                    np.array(logits[:, -1, :], dtype=np.float32, copy=True),
+                    temperature=temperature,
+                    top_k=top_k,
+                )
+                tokens.append(next_token)
+
+                if next_token == tokenizer.eos_id:
+                    break
+
+                if len(tokens) >= CONTEXT:
+                    break
+
+        # Decode output
+        generated_text = tokenizer.decode(tokens)
+        return generated_text
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# ============================================================
+# Gradio UI
+# ============================================================
+
+def main():
+    with gr.Blocks(title="Veylon Alpha") as demo:
+        gr.Markdown("""
+        # 🚀 Veylon Alpha Preview - 10M LLM
+        # Made by Arush Kumar
+        A student dev! 
+        A small transformer model trained on clean data.
+        Enter a prompt and watch it generate text.
+        
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Once upon a time",
+                    lines=3,
+                    value="Once upon a time"
+                )
 
-    def __init__(self, d_model, num_experts=8, top_k=2, hidden_mult=3.5, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model_arg = d_model
-        self.num_experts = num_experts
-        self.top_k = top_k
-        self.hidden_mult = hidden_mult
-        self.experts = [SwiGLUFFN(d_model, hidden_mult) for _ in range(num_experts)]
-        self.router = layers.Dense(num_experts, use_bias=False)
-
-    def _load_balancing_loss(self, router_logits, top_idx):
-        # Standard Switch-Transformer aux loss.  `f` uses the argmax routing
-        # decision (top_idx[..., 0]); `p` uses the router softmax mean.
-        router_probs = ops.softmax(router_logits, axis=-1)
-        mask = ops.one_hot(top_idx[..., 0], self.num_experts)        # [..., num_experts]
-        f = ops.mean(ops.cast(mask, 'float32'), axis=tuple(range(mask.ndim - 1)))
-        p = ops.mean(router_probs, axis=tuple(range(router_probs.ndim - 1)))
-        return self.num_experts * ops.sum(f * p)
-
-    def call(self, x, training=False):
-        # x: [B, S, D]
-        B = ops.shape(x)[0]
-        S = ops.shape(x)[1]
-        D = x.shape[-1]
-
-        router_logits = self.router(x)                              # [B, S, E]
-        top_logits, top_idx = ops.top_k(router_logits, self.top_k)  # [B, S, k]
-        top_weights = ops.softmax(top_logits, axis=-1)              # [B, S, k]
-
-        if training:
-            self.add_loss(self._load_balancing_loss(router_logits, top_idx))
-
-        out = ops.zeros_like(x)
-        for kk in range(self.top_k):
-            idx_k = top_idx[..., kk]                                 # [B, S]
-            w_k = top_weights[..., kk][..., None]                    # [B, S, 1]
-            for e in range(self.num_experts):
-                mask_e = ops.cast(idx_k == e, x.dtype)[..., None]    # [B, S, 1]
-                # mask_e is always a tensor; skip the useless ops.is_tensor guard
-                # (it never triggered — removed FIX BUG11)
-                # Only run the expert where it is actually selected.
-                x_e = x * mask_e
-                expert_out = self.experts[e](x_e, training=training)
-                out = out + expert_out * w_k * mask_e
-        return out
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'd_model': self.d_model_arg, 'num_experts': self.num_experts,
-            'top_k': self.top_k, 'hidden_mult': self.hidden_mult,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class MLAttention(layers.Layer):
-    def __init__(self, d_model, n_heads, d_latent, max_seq_len, num_kv_heads=2,
-                 swa_window=1024, attn_dropout=0.0, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        if d_model % n_heads != 0:
-            raise ValueError('d_model must be divisible by n_heads')
-        if n_heads % num_kv_heads != 0:
-            raise ValueError('n_heads must be divisible by num_kv_heads')
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.num_kv_heads = num_kv_heads
-        self.group_size = n_heads // num_kv_heads
-        self.d_head = d_model // n_heads
-        self.d_latent = d_latent
-        self.max_seq_len = max_seq_len
-        self.swa_window = swa_window
-        self.gen_headroom = int(gen_headroom)
-        self.dropout = layers.Dropout(attn_dropout)
-        self.rope = RotaryEmbedding(max_seq_len, d_model // n_heads, gen_headroom=int(gen_headroom))  # FIX BUG10: must be in __init__ for Keras tracking
-
-    def build(self, input_shape):
-        self.W_qc = self.add_weight(shape=(self.d_model, self.d_model + self.d_latent), initializer='glorot_uniform', name='W_qc')
-        self.W_kv = self.add_weight(shape=(self.d_latent, self.num_kv_heads * 2 * self.d_head), initializer='glorot_uniform', name='W_kv')
-        self.W_o = self.add_weight(shape=(self.d_model, self.d_model), initializer='glorot_uniform', name='Wo')
-
-    def _project_kv(self, c):
-        kv = ops.matmul(c, self.W_kv)
-        return ops.split(kv, 2, axis=-1)
-
-    def call(self, x, training=False):
-        B = ops.shape(x)[0]
-        S = ops.shape(x)[1]
-        qc = ops.matmul(x, self.W_qc)
-        q_proj, c = ops.split(qc, [self.d_model], axis=-1)
-        q = ops.reshape(q_proj, (B, S, self.n_heads, self.d_head))
-        q = self.rope(q, offset=0)
-        q = ops.transpose(q, (0, 2, 1, 3))
-        k, v = self._project_kv(c)
-        k = ops.reshape(k, (B, S, self.num_kv_heads, self.d_head))
-        v = ops.reshape(v, (B, S, self.num_kv_heads, self.d_head))
-        k = self.rope(k, offset=0)
-        k = ops.transpose(k, (0, 2, 1, 3))
-        v = ops.transpose(v, (0, 2, 1, 3))
-
-        # Attention remat is REQUIRED during training: it is the per-block
-        # jax.checkpoint(process_block) that stops XLA from staging the hidden
-        # [S, B, Hkv, W, D] buffer (≈8.6 GB here) inside the fori_loop.  Block-
-        # level remat does NOT prevent that forward-time allocation, so this
-        # must follow `training` regardless of config.USE_REMAT.
-        out = flash_splash_attention(
-            q, k, v,
-            window_size=min(self.swa_window, self.max_seq_len),
-            backend=jax.default_backend(),
-            use_gqa=True,
-            use_remat=training,
-        )
-        out = ops.transpose(out, (0, 2, 1, 3))
-        out = ops.reshape(out, (B, S, self.d_model))
-        out = self.dropout(out, training=training)
-        return ops.matmul(out, self.W_o)
-
-    def generate_step(self, x, cache_k=None, cache_v=None, cache_pos=0):
-        B = ops.shape(x)[0]
-        S = ops.shape(x)[1]
-
-        qc = ops.matmul(x, self.W_qc)
-        q_proj, c = ops.split(qc, [self.d_model], axis=-1)
-
-        q = ops.reshape(q_proj, (B, S, self.n_heads, self.d_head))
-        q = self.rope(q, offset=cache_pos)
-        q = ops.transpose(q, (0, 2, 1, 3))
-
-        k, v = self._project_kv(c)
-        k = ops.reshape(k, (B, S, self.num_kv_heads, self.d_head))
-        v = ops.reshape(v, (B, S, self.num_kv_heads, self.d_head))
-        k = self.rope(k, offset=cache_pos)
-        k = ops.transpose(k, (0, 2, 1, 3))
-        v = ops.transpose(v, (0, 2, 1, 3))
-
-        if cache_k is None:
-            # Prefill: no remat (inference only), no window trim needed.
-            out = flash_splash_attention(
-                q, k, v,
-                window_size=min(self.swa_window, self.max_seq_len),
-                backend=jax.default_backend(),
-                use_gqa=True,
-                use_remat=False,
-            )
-            new_k = k[:, :, -self.swa_window:, :]
-            new_v = v[:, :, -self.swa_window:, :]
-        else:
-            if S != 1:
-                raise ValueError(f'generate_step with cache expects S=1, got S={S}')
-
-            k = ops.concatenate([cache_k, k], axis=2)
-            v = ops.concatenate([cache_v, v], axis=2)
-
-            k = k[:, :, -self.swa_window:, :]
-            v = v[:, :, -self.swa_window:, :]
-
-            out = decode_swa(q, k, v)
-            new_k = k
-            new_v = v
-
-        out = ops.transpose(out, (0, 2, 1, 3))
-        out = ops.reshape(out, (B, S, self.d_model))
-        out = ops.matmul(out, self.W_o)
-        return out, new_k, new_v
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'd_model': self.d_model, 'n_heads': self.n_heads,
-            'num_kv_heads': self.num_kv_heads, 'd_latent': self.d_latent,
-            'max_seq_len': self.max_seq_len, 'swa_window': self.swa_window,
-            'attn_dropout': self.dropout.rate, 'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class TransformerBlock(layers.Layer):
-    def __init__(self, d_model, n_heads, d_latent, ffn_layer, max_seq_len,
-                 num_kv_heads=2, swa_window=1024, use_remat=True, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.d_latent = d_latent
-        self.max_seq_len = max_seq_len
-        self.num_kv_heads = num_kv_heads
-        self.swa_window = swa_window
-        self.use_remat = use_remat
-        self.gen_headroom = int(gen_headroom)
-        self.ffn = keras.saving.deserialize_keras_object(ffn_layer) if isinstance(ffn_layer, dict) else ffn_layer
-        self.norm1 = RMSNorm()
-        self.norm2 = RMSNorm()
-        self.attn = MLAttention(
-            d_model, n_heads, d_latent, max_seq_len,
-            num_kv_heads=num_kv_heads, swa_window=swa_window, gen_headroom=self.gen_headroom,
+                with gr.Row():
+                    max_tokens = gr.Slider(
+                        label="Max tokens",
+                        minimum=10,
+                        maximum=256,
+                        value=64,
+                        step=10,
+                    )
+                    temperature = gr.Slider(
+                        label="Temperature",
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=0.8,
+                        step=0.1,
+                    )
+                    top_k = gr.Slider(
+                        label="Top-K",
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=1,
+                    )
+
+                generate_btn = gr.Button("Generate", variant="primary", size="lg")
+
+            with gr.Column(scale=1):
+                info = gr.Markdown(f"""
+                **Model Info**
+                
+                - Parameters: {model.count_params():,}
+                - Context: {CONTEXT} tokens
+                - Vocab: {vocab_size}
+                - Architecture: Transformer + GQA
+                
+                **Tips**
+                - Higher temp = more creative
+                - Lower temp = more deterministic
+                - Top-K = diversity control
+                """)
+
+        output = gr.Textbox(
+            label="Generated Output",
+            lines=8,
+            interactive=False
         )
 
-    def call(self, x, training=False):
-        # Block-level activation/gradient checkpointing (optional, via
-        # config.USE_REMAT).  NOTE: this is independent of attention's own
-        # remat, which MUST stay on whenever training=True — that's what stops
-        # XLA staging the hidden [S,B,Hkv,W,D] buffer (the OOM we hit).
-        if training and self.use_remat:
-            def _fwd(x_in):
-                attn_out = self.attn(self.norm1(x_in), training=True)
-                ffn_out = self.ffn(self.norm2(x_in), training=True)
-                return attn_out + ffn_out
-            return x + jax.checkpoint(_fwd)(x)
-        a = self.attn(self.norm1(x), training=training)
-        f = self.ffn(self.norm2(x), training=training)
-        return x + a + f
-
-    def generate_step(self, x, cache_k=None, cache_v=None, cache_pos=0):
-        attn_out, nck, ncv = self.attn.generate_step(
-            self.norm1(x), cache_k=cache_k, cache_v=cache_v, cache_pos=cache_pos,
+        # Connect
+        generate_btn.click(
+            fn=generate,
+            inputs=[prompt, max_tokens, temperature, top_k],
+            outputs=output,
+            api_name="generate"
         )
-        x = x + attn_out
-        x = x + self.ffn(self.norm2(x), training=False)
-        return x, nck, ncv
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'd_model': self.d_model, 'n_heads': self.n_heads, 'd_latent': self.d_latent,
-            'ffn_layer': keras.saving.serialize_keras_object(self.ffn),
-            'max_seq_len': self.max_seq_len, 'num_kv_heads': self.num_kv_heads,
-            'swa_window': self.swa_window, 'use_remat': self.use_remat,
-            'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-@keras.saving.register_keras_serializable()
-class VeylonModel(keras.Model):
-    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_latent, ffn_mult,
-                 max_seq_len, use_moe=False, moe_num_experts=8, moe_top_k=2,
-                 num_kv_heads=2, swa_window=1024, use_remat=True, gen_headroom=0, **kwargs):
-        super().__init__(**kwargs)
-        self.vocab_size = vocab_size
-        self.d_model = d_model
-        self.n_layers = n_layers
-        self.n_heads = n_heads
-        self.d_latent = d_latent
-        self.ffn_mult = ffn_mult
-        self.max_seq_len = max_seq_len
-        self.use_moe = use_moe
-        self.moe_num_experts = moe_num_experts
-        self.moe_top_k = moe_top_k
-        self.num_kv_heads = num_kv_heads
-        self.swa_window = swa_window
-        self.use_remat = use_remat
-        self.gen_headroom = int(gen_headroom)
-        self.embedding = layers.Embedding(vocab_size, d_model, name='token_embedding')
-        self.blocks = []
-        for i in range(n_layers):
-            ffn = (MoE_FFN(d_model, moe_num_experts, moe_top_k, ffn_mult)
-                   if use_moe else SwiGLUFFN(d_model, ffn_mult))
-            self.blocks.append(TransformerBlock(
-                d_model, n_heads, d_latent, ffn, max_seq_len,
-                num_kv_heads=num_kv_heads, swa_window=swa_window,
-                use_remat=use_remat, gen_headroom=self.gen_headroom, name=f'block_{i}',
-            ))
-        self.norm = RMSNorm()
-
-    def call(self, inputs, training=False):
-        x = self.embedding(inputs)
-        for block in self.blocks:
-            x = block(x, training=training)
-        x = self.norm(x)
-        embedding_weights = self.embedding.embeddings
-        logits = ops.matmul(x, ops.transpose(embedding_weights))
-        return ops.cast(logits, 'float32')
-
-    def generate_step(self, inputs, cache_k=None, cache_v=None, cache_pos=0):
-        x = self.embedding(inputs)
-        new_cache_k = []
-        new_cache_v = []
-
-        if cache_k is None:
-            cache_k = [None] * len(self.blocks)
-            cache_v = [None] * len(self.blocks)
-
-        for i, block in enumerate(self.blocks):
-            x, nck, ncv = block.generate_step(
-                x, cache_k=cache_k[i], cache_v=cache_v[i], cache_pos=cache_pos,
-            )
-            new_cache_k.append(nck)
-            new_cache_v.append(ncv)
-
-        x = self.norm(x)
-        embedding_weights = self.embedding.embeddings
-        logits = ops.matmul(x, ops.transpose(embedding_weights))
-        logits = ops.cast(logits, 'float32')
-        return logits, new_cache_k, new_cache_v
-
-    def get_config(self):
-        cfg = super().get_config()
-        cfg.update({
-            'vocab_size': self.vocab_size, 'd_model': self.d_model,
-            'n_layers': self.n_layers, 'n_heads': self.n_heads,
-            'd_latent': self.d_latent, 'ffn_mult': self.ffn_mult,
-            'max_seq_len': self.max_seq_len, 'use_moe': self.use_moe,
-            'moe_num_experts': self.moe_num_experts, 'moe_top_k': self.moe_top_k,
-            'num_kv_heads': self.num_kv_heads, 'swa_window': self.swa_window,
-            'use_remat': self.use_remat, 'gen_headroom': self.gen_headroom,
-        })
-        return cfg
-
-
-# ─────────────────────────────────────────────────────────────────────────────
-# Factory
-# ─────────────────────────────────────────────────────────────────────────────
-
-def create_llm(
-    vocab_size=Vocab_size, d_model=D_MODEL, n_layers=numberoflayers, n_heads=numberofheads,
-    d_latent=d_Latent, ffn_mult=ffn_mult, max_seq_len=CONTEXT, use_moe=use_moe,
-    moe_num_experts=moe_num_experts, moe_top_k=moe_top_k, num_kv_heads=num_kv_heads,
-    swa_window=swa_window, use_remat=USE_REMAT, gen_headroom=MAX_GEN_TOKENS,
-):
-    """
-    Build a VeylonModel.
 
-    `gen_headroom` sizes the RoPE table to (max_seq_len + gen_headroom) so that
-    a full-length prompt can still generate `gen_headroom` tokens during
-    inference without "RoPE table too small" errors.  Defaults to MAX_GEN_TOKENS.
-    """
-    return VeylonModel(
-        vocab_size=vocab_size, d_model=d_model, n_layers=n_layers, n_heads=n_heads,
-        d_latent=d_latent, ffn_mult=ffn_mult, max_seq_len=max_seq_len, use_moe=use_moe,
-        moe_num_experts=moe_num_experts, moe_top_k=moe_top_k, num_kv_heads=num_kv_heads,
-        swa_window=swa_window, use_remat=use_remat, gen_headroom=gen_headroom,
-    )
+    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
+
+if __name__ == "__main__":
+    main()