Update app.py

app.py

@@ -3,11 +3,17 @@ from __future__ import annotations
 import os
 os.environ["KERAS_BACKEND"] = "jax"
 
+import time
+from pathlib import Path
+from functools import partial
+
 import numpy as np
 import jax
+import jax.numpy as jnp
 import keras
 import gradio as gr
-from pathlib import Path
+from fastapi import FastAPI
+import uvicorn
 
 from veylon_model import create_llm
 from tokenizer import TokenizerWrapper
@@ -24,7 +30,25 @@ from config import (
 )
 
 # ============================================================
-# Initialize (runs once)
+# Config
+# ============================================================
+
+WEIGHTS_PATH = "veylon_final.weights.h5"
+MODEL_NAME = "Arya"
+MODEL_TAGLINE = "27M-param transformer, made by Arush Kumar"
+
+# Chat templating for a BASE (non-instruction-tuned) language model. If
+# Arya/Veylon was trained on plain prose (TinyStories/WikiText/etc.) rather
+# than dialogue-formatted data, this template is best-effort — multi-turn
+# coherence depends entirely on what the training corpus looked like, not
+# on anything this app can fix. Swap this format if you fine-tune on a
+# specific chat template later.
+USER_TAG = "User:"
+ASSISTANT_TAG = "Assistant:"
+
+
+# ============================================================
+# Initialize (runs once, at import time)
 # ============================================================
 
 keras.mixed_precision.set_global_policy("mixed_bfloat16")
@@ -32,14 +56,12 @@ 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")
+print(f"Tokenizer loaded: {tokenizer.vocab_size} vocab")
 
-# Build model
 print("Building model...")
 model = create_llm(
     vocab_size=vocab_size,
@@ -54,21 +76,59 @@ model = create_llm(
     swa_window=swa_window,
 )
 
-# Warmup
 dummy = np.zeros((1, CONTEXT), dtype=np.int32)
 _ = model(dummy, training=False)
-print("✓ Model built successfully")
+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")
+    print("Weights loaded successfully")
 else:
     print(f"WARNING: {WEIGHTS_PATH} not found. Using untrained model.")
 
-print(f"✓ Model params: {model.count_params():,}\n")
+print(f"Model params: {model.count_params():,}\n")
+
+
+# ============================================================
+# JIT-compiled generation steps — the real speed lever here.
+# ============================================================
+# The original app called model.generate_step() as plain, untraced Python
+# on every single token — meaning JAX re-dispatches and re-optimizes the
+# whole graph every call instead of running a compiled executable. This is
+# the single biggest speed cost in the original code, and (unlike swapping
+# inference engines) is something we can fix directly: jax.jit traces the
+# step function once per unique input shape/dtype, caches the compiled
+# XLA executable, and every subsequent call re-uses it. Since prefill
+# always sees shape (1, CONTEXT) and decode always sees shape (1, 1),
+# there are exactly two shapes to compile — both get cached on first use,
+# after which every call is a compiled dispatch instead of a Python/XLA
+# re-trace.
+#
+# NOTE: this assumes model.generate_step is a pure function of its
+# arguments (no hidden Python-side mutation of the cache/model state
+# outside of what it returns). That's true for essentially all KV-cache
+# transformer implementations, but if your model.generate_step does
+# something stateful and doesn't work under jit, set JIT_GENERATE=False
+# below and it'll fall back to the original eager behavior with zero
+# other changes.
+
+JIT_GENERATE = True
+
+if JIT_GENERATE:
+    _generate_step_jit = jax.jit(model.generate_step)
+else:
+    _generate_step_jit = model.generate_step
+
+
+def _run_generate_step(tokens_arr, cache_k, cache_v, cache_pos):
+    return _generate_step_jit(
+        tokens_arr,
+        cache_k=cache_k,
+        cache_v=cache_v,
+        cache_pos=cache_pos,
+    )
+
 
 # ============================================================
 # Sampling
@@ -100,48 +160,71 @@ def sample_from_logits(
 
     return int(np.random.choice(len(probs), p=probs))
 
+
 # ============================================================
-# Generation function
+# Prompt construction for chat-style multi-turn context
 # ============================================================
 
-def generate(
-    prompt: str,
-    max_new_tokens: int = 64,
-    temperature: float = 0.8,
-    top_k: int = 50,
-) -> str:
+def build_prompt(message: str, history: list[dict]) -> str:
+    """
+    Turns Gradio chat history (list of {"role", "content"} dicts, the
+    'messages' format) into a single text block for a base LM, truncated
+    to fit CONTEXT. Oldest turns are dropped first if it doesn't fit —
+    we build from the end backwards.
+    """
+    turns = []
+    for turn in history:
+        role = turn.get("role")
+        content = turn.get("content", "")
+        tag = USER_TAG if role == "user" else ASSISTANT_TAG
+        turns.append(f"{tag} {content}")
+    turns.append(f"{USER_TAG} {message}")
+    turns.append(ASSISTANT_TAG)
+
+    # Build from the most recent turn backwards until we run out of
+    # context budget, so we never silently truncate the *current* message.
+    budget_chars = CONTEXT * 4  # rough chars-per-token heuristic pre-tokenize
+    kept = []
+    running_len = 0
+    for t in reversed(turns):
+        running_len += len(t) + 1
+        if running_len > budget_chars and kept:
+            break
+        kept.append(t)
+    kept.reverse()
+    return "\n".join(kept)
+
+
+# ============================================================
+# Generation (streaming generator for gr.ChatInterface)
+# ============================================================
+
+def generate_stream(
+    message: str,
+    history: list[dict],
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+):
     """
-    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
+    Streaming generator: yields the partial assistant response as tokens
+    are produced, so the chat UI fills in live instead of blocking on the
+    whole generation.
     """
     try:
-        # Encode prompt
-        tokens = tokenizer.encode(
-            prompt,
-            add_bos=True,
-            add_eos=False,
-        )
+        prompt = build_prompt(message, history)
 
+        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:]
 
-        # Prefill phase
+        prompt_len = len(tokens)
         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,
+
+        t0 = time.time()
+        logits, cache_k, cache_v = _run_generate_step(
+            prompt_ids, cache_k=None, cache_v=None, cache_pos=0,
         )
 
         next_token = sample_from_logits(
@@ -149,22 +232,20 @@ def generate(
             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])
+        generated_ids = [next_token]
+        partial_text = tokenizer.decode(generated_ids)
+        yield partial_text
+
+        if next_token != tokenizer.eos_id and prompt_len < CONTEXT:
+            cache_pos = prompt_len
 
-            for _ in range(max_new_tokens - 1):
+            for _ in range(int(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,
+                logits, cache_k, cache_v = _run_generate_step(
+                    next_input, cache_k=cache_k, cache_v=cache_v, cache_pos=cache_pos,
                 )
-
                 cache_pos += 1
 
                 next_token = sample_from_logits(
@@ -172,100 +253,221 @@ def generate(
                     temperature=temperature,
                     top_k=top_k,
                 )
-                tokens.append(next_token)
 
                 if next_token == tokenizer.eos_id:
                     break
 
-                if len(tokens) >= CONTEXT:
+                generated_ids.append(next_token)
+                partial_text = tokenizer.decode(generated_ids)
+                yield partial_text
+
+                if prompt_len + len(generated_ids) >= CONTEXT:
                     break
 
-        # Decode output
-        generated_text = tokenizer.decode(tokens)
-        return generated_text
+        elapsed = time.time() - t0
+        n_tok = len(generated_ids)
+        tok_per_sec = n_tok / elapsed if elapsed > 0 else 0.0
+        print(f"Generated {n_tok} tokens in {elapsed:.2f}s ({tok_per_sec:.1f} tok/s)")
 
     except Exception as e:
-        return f"Error: {str(e)}"
+        yield f"⚠ Error: {str(e)}"
+
 
 # ============================================================
-# Gradio UI
+# Styling — dark terminal / glassmorphism, JetBrains Mono, neon accent
 # ============================================================
 
-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.
-        
+CUSTOM_CSS = """
+@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;700&display=swap');
+
+:root {
+    --veylon-bg: #0a0e14;
+    --veylon-panel: rgba(20, 26, 38, 0.65);
+    --veylon-border: rgba(0, 255, 200, 0.18);
+    --veylon-accent: #00ffc8;
+    --veylon-accent-dim: #00b894;
+    --veylon-text: #e6f1ff;
+    --veylon-text-dim: #7d8ba1;
+}
+
+* {
+    font-family: 'JetBrains Mono', ui-monospace, monospace !important;
+}
+
+body, .gradio-container {
+    background: radial-gradient(circle at 20% 10%, #101826 0%, #0a0e14 60%) !important;
+    color: var(--veylon-text) !important;
+}
+
+#veylon-header {
+    text-align: center;
+    padding: 18px 12px 6px 12px;
+}
+
+#veylon-header h1 {
+    font-size: 1.8rem;
+    font-weight: 700;
+    letter-spacing: 0.5px;
+    background: linear-gradient(90deg, var(--veylon-accent), #4dd0ff);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    margin-bottom: 2px;
+}
+
+#veylon-header p {
+    color: var(--veylon-text-dim);
+    font-size: 0.9rem;
+    margin: 0;
+}
+
+.gr-panel, .block, .form {
+    background: var(--veylon-panel) !important;
+    border: 1px solid var(--veylon-border) !important;
+    border-radius: 14px !important;
+    backdrop-filter: blur(10px);
+}
+
+#chatbot {
+    border-radius: 16px !important;
+    border: 1px solid var(--veylon-border) !important;
+    background: var(--veylon-panel) !important;
+    box-shadow: 0 0 24px rgba(0, 255, 200, 0.05);
+}
+
+.message.user {
+    background: rgba(0, 255, 200, 0.08) !important;
+    border: 1px solid var(--veylon-border) !important;
+}
+
+.message.bot {
+    background: rgba(255, 255, 255, 0.03) !important;
+    border: 1px solid rgba(255, 255, 255, 0.06) !important;
+}
+
+button.primary, .gr-button-primary {
+    background: linear-gradient(90deg, var(--veylon-accent-dim), var(--veylon-accent)) !important;
+    color: #04120e !important;
+    border: none !important;
+    font-weight: 700 !important;
+    box-shadow: 0 0 16px rgba(0, 255, 200, 0.25);
+    transition: transform 0.12s ease, box-shadow 0.12s ease;
+}
+
+button.primary:hover, .gr-button-primary:hover {
+    transform: translateY(-1px);
+    box-shadow: 0 0 24px rgba(0, 255, 200, 0.4);
+}
+
+input, textarea, .gr-box {
+    background: rgba(255, 255, 255, 0.03) !important;
+    color: var(--veylon-text) !important;
+    border: 1px solid var(--veylon-border) !important;
+    border-radius: 10px !important;
+}
+
+label span {
+    color: var(--veylon-text-dim) !important;
+    font-size: 0.78rem !important;
+    text-transform: uppercase;
+    letter-spacing: 0.06em;
+}
+
+#stats-badge {
+    text-align: center;
+    color: var(--veylon-text-dim);
+    font-size: 0.78rem;
+    padding-top: 4px;
+}
+
+footer { display: none !important; }
+"""
+
+THEME = gr.themes.Base(
+    primary_hue="teal",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
+).set(
+    body_background_fill="*neutral_950",
+    block_background_fill="*neutral_900",
+    border_color_primary="*primary_700",
+)
+
+
+# ============================================================
+# Build the Gradio Blocks app
+# ============================================================
+
+def build_demo() -> gr.Blocks:
+    with gr.Blocks(title=MODEL_NAME, theme=THEME, css=CUSTOM_CSS) as demo:
+        gr.HTML(f"""
+        <div id="veylon-header">
+            <h1>◆ {MODEL_NAME}</h1>
+            <p>{MODEL_TAGLINE} · {model.count_params():,} params · {CONTEXT} ctx</p>
+        </div>
         """)
 
         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"
-                )
-
-                with gr.Row():
-                    max_tokens = gr.Slider(
-                        label="Max tokens",
-                        minimum=64,
-                        maximum=768,
-                        value=128,
-                        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
+            max_tokens = gr.Slider(
+                label="Max new tokens", minimum=10, maximum=768, 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,
+            )
+
+        gr.ChatInterface(
+            fn=generate_stream,
+            type="messages",
+            additional_inputs=[max_tokens, temperature, top_k],
+            chatbot=gr.Chatbot(elem_id="chatbot", type="messages", height=480),
+            examples=[
+                ["Once upon a time,"],
+                ["The scientist looked at the data and"],
+                ["In a small village, there lived"],
+            ],
+            title=None,
         )
 
-        # Connect
-        generate_btn.click(
-            fn=generate,
-            inputs=[prompt, max_tokens, temperature, top_k],
-            outputs=output,
-            api_name="generate"
-        )
+        gr.HTML(f"""
+        <div id="stats-badge">
+            base language model · not instruction-tuned · coherence across
+            turns depends on training data format
+        </div>
+        """)
+
+    return demo
+
+
+demo = build_demo()
+demo.queue()  # required for streaming generators to work over the wire
+
+
+# ============================================================
+# FastAPI mount — real production entrypoint instead of demo.launch()
+# ============================================================
+
+app = FastAPI(title=f"{MODEL_NAME} API")
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok",
+        "model": MODEL_NAME,
+        "params": model.count_params(),
+        "context": CONTEXT,
+        "backend": keras.backend.backend(),
+        "jax_devices": [str(d) for d in jax.devices()],
+    }
+
+
+app = gr.mount_gradio_app(app, demo, path="/")
 
-    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)
 
 if __name__ == "__main__":
-    main()
+    # Run via uvicorn directly rather than demo.launch(): gives you
+    # multi-worker / reload / proper ASGI server config for actual
+    # deployment, plus the /health route above for load-balancer checks.
+    uvicorn.run(app, host="0.0.0.0", port=7860)