Update app.py

app.py

@@ -1,558 +1,655 @@
+"""
+app.py — Pragya by Arush Kumar · Hugging Face Spaces (PyTorch backend)
+
+Expects alongside app.py:
+    model.py       – GPT / GPTConfig (training codebase, unchanged)
+    inference.py   – generate(), load_model(), sampling helpers
+    tokenizer.py   – TokenizerWrapper
+    tokenizer.model – SentencePiece vocab file
+    checkpoint.pt  – trained checkpoint (from train.py)
+
+Environment variables (Space Settings → Variables):
+    CHECKPOINT_PATH   path to local checkpoint  (default ./checkpoint.pt)
+    CHECKPOINT_REPO   HF Hub repo id            (e.g. yourname/pragya)
+    CHECKPOINT_FILE   filename inside Hub repo  (default checkpoint.pt)
+    TOKENIZER_PATH    path to tokenizer.model   (default ./tokenizer.model)
+    MAX_TOKENS_CAP    hard cap on slider         (default 512)
+    DEFAULT_MAX_TOKENS default slider value      (default 128)
+"""
+
 from __future__ import annotations
 
+import contextlib
 import os
-os.environ["KERAS_BACKEND"] = "jax"
-
-import numpy as np
-import jax
-import keras
-import gradio as gr
+import sys
 import time
+import threading
 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")
+import torch
+import gradio as gr
 
-print(f"Backend: {keras.backend.backend()}")
-print(f"JAX devices: {jax.devices()}")
+_HERE = Path(__file__).resolve().parent
+if str(_HERE) not in sys.path:
+    sys.path.insert(0, str(_HERE))
 
-# 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,
+from inference import (         # noqa: E402
+    load_model, pick_device, _resolve_dtype,
+    _apply_repetition_penalty, _block_repeated_ngrams,
+    _top_k_filter, _top_p_filter,
 )
+from tokenizer import TokenizerWrapper  # noqa: E402
+
+# ── env config ───────────────────────────────────────────────────────────────
+_CHECKPOINT_LOCAL   = os.environ.get("CHECKPOINT_PATH",    str(_HERE / "checkpoint.pt"))
+_HF_REPO            = os.environ.get("CHECKPOINT_REPO",    "")
+_HF_FILE            = os.environ.get("CHECKPOINT_FILE",    "pragya.pt")
+_TOKENIZER          = os.environ.get("TOKENIZER_PATH",     str(_HERE / "tokenizer.model"))
+_MAX_TOKENS_CAP     = int(os.environ.get("MAX_TOKENS_CAP",      "512"))
+_DEFAULT_TOKENS     = int(os.environ.get("DEFAULT_MAX_TOKENS",  "128"))
+
+
+def _resolve_checkpoint() -> str:
+    if _HF_REPO and _HF_FILE:
+        try:
+            from huggingface_hub import hf_hub_download
+            print(f"[App] Pulling {_HF_REPO}/{_HF_FILE} from HF Hub …", flush=True)
+            return hf_hub_download(repo_id=_HF_REPO, filename=_HF_FILE)
+        except Exception as exc:
+            print(f"[App] Hub download failed ({exc}), trying local.", flush=True)
+    if Path(_CHECKPOINT_LOCAL).exists():
+        return _CHECKPOINT_LOCAL
+    raise FileNotFoundError(
+        f"No checkpoint at {_CHECKPOINT_LOCAL}. "
+        "Set CHECKPOINT_REPO + CHECKPOINT_FILE to pull from HF Hub."
+    )
 
-# 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_row: np.ndarray,
-    temperature: float = 0.7,
-    min_p: float = 0.05,
-) -> int:
-    """
-    Min-P sampling (replaces top-k). Top-k/top-p keep a FIXED-size or
-    fixed-cumulative-mass candidate pool open regardless of how confident
-    the model actually is -- on a small/uncertain model, that means when
-    it doesn't know what comes next, top-k still hands it 50 candidates
-    to choose from, many of them garbage. Min-P instead sets a threshold
-    RELATIVE to the top token's probability: confident predictions
-    collapse the pool to 1-2 tokens, uncertain predictions keep it wide.
-    Expects and returns a 1D array (single vocab-length row), not (1, vocab).
-    """
-    logits_row = np.array(logits_row, dtype=np.float32, copy=True)
 
-    if temperature > 0:
-        logits_row = logits_row / float(max(temperature, 1e-8))
+# ── one-time model init ───────────────────────────────────────────────────────
+print("[App] Initialising …", flush=True)
 
-    row = logits_row - np.max(logits_row)  # numerical stability
-    probs = np.exp(row)
-    probs = probs / probs.sum()
+_device         = pick_device(None)
+_cast_dtype     = _resolve_dtype("auto", _device)
+_ckpt_path      = _resolve_checkpoint()
+_model, _cfg    = load_model(_ckpt_path, _device)
+_model.eval()
+_tokenizer      = TokenizerWrapper(_TOKENIZER)
 
-    if min_p > 0:
-        threshold = min_p * probs.max()
-        mask = probs >= threshold
-        if not mask.any():  # degenerate guard -- never zero out everything
-            mask[np.argmax(probs)] = True
-        probs = np.where(mask, probs, 0.0)
-        probs = probs / probs.sum()
+_n_params_str   = f"{_model.get_num_params() / 1e6:.1f}M"
+_ctx_len        = _cfg.block_size
+_pattern        = getattr(_cfg, "pattern", "dense").upper()
+_n_layer        = getattr(_cfg, "n_layer",  "?")
+_n_head         = getattr(_cfg, "n_head",   "?")
+_kv_heads       = getattr(_cfg, "num_kv_heads", _n_head)
+_gen_lock       = threading.Lock()
 
-    return int(np.random.choice(len(probs), p=probs))
+print(f"[App] Ready — {_n_params_str} params  ctx={_ctx_len}  device={_device}", flush=True)
 
 
-def apply_repetition_controls(
-    logits_row: np.ndarray,
-    generated_ids: list,
-    no_repeat_ngram_size: int = 3,
-    repetition_penalty: float = 1.2,
-) -> np.ndarray:
-    """
-    Applied BEFORE sample_from_logits, on the model's own newly-generated
-    tokens only (never the user's prompt -- penalizing someone for the
-    model reusing words from THEIR prompt would be wrong, this is about
-    stopping the model's own loops).
-
-    1. repetition_penalty (soft): divides a positive logit for an
-       already-seen token, MULTIPLIES a negative one -- always pushes the
-       token down, regardless of logit sign (a naive "divide by penalty"
-       for a negative logit would actually push it UP, backwards).
-    2. no_repeat_ngram_size (hard): if continuing with a candidate token
-       would recreate an n-gram that's already appeared in this
-       generation, that candidate is set to -inf. This is what actually
-       kills infinite loops -- it's a hard constraint, not a nudge.
-    """
-    logits_row = logits_row.copy()
-
-    if repetition_penalty and repetition_penalty != 1.0 and generated_ids:
-        for tok_id in set(generated_ids):
-            if logits_row[tok_id] > 0:
-                logits_row[tok_id] /= repetition_penalty
-            else:
-                logits_row[tok_id] *= repetition_penalty
-
-    n = no_repeat_ngram_size
-    if n and n > 0 and len(generated_ids) >= n - 1:
-        prefix = tuple(generated_ids[-(n - 1):]) if n > 1 else tuple()
-        banned = set()
-        for i in range(len(generated_ids) - n + 1):
-            if tuple(generated_ids[i:i + n - 1]) == prefix:
-                banned.add(generated_ids[i + n - 1])
-        for tok_id in banned:
-            logits_row[tok_id] = -np.inf
-
-    return logits_row
-
-
-# Few-shot "anchoring": Veylon is a CONTINUATION-only base model, never
-# instruction-tuned -- it has no concept of "Input:"/"Output:" turns, so
-# priming it with instruction-style examples (as generic small-LM advice
-# often suggests) would push it OUT of its training distribution, not
-# into it. What anchoring means correctly here is different: prepend a
-# short, clean, well-formed piece of narrative prose so the model's
-# recent-attention window is dominated by coherent English *before*
-# generation starts, nudging it away from degenerating into word salad.
-# This consumes some of the context budget, so it's short and toggleable.
-ANCHOR_TEXT = (
-    "The old lighthouse stood at the edge of the cliff, its light "
-    "sweeping slowly across the dark water. Every night, the keeper "
-    "climbed the narrow stairs to check the lamp, and every night the "
-    "sea answered with the steady sound of waves against the rocks.\n\n"
+# ── generation ────────────────────────────────────────────────────────────────
+_ANCHOR = (
+    "The old lighthouse stood at the edge of the cliff, its beam sweeping "
+    "slowly across the dark water. Every night the keeper climbed the narrow "
+    "stairs to check the lamp, and every night the sea answered with the steady "
+    "sound of waves breaking against the rocks far below.\n\n"
 )
 
-# ============================================================
-# Streaming generation with live telemetry
-# ============================================================
+def _cast_ctx():
+    if _cast_dtype is not None and _device.type in ("cuda", "cpu"):
+        return torch.autocast(device_type=_device.type, dtype=_cast_dtype)
+    return contextlib.nullcontext()
+
 
 def respond(
     message: str,
     history: list,
-    max_new_tokens: int = 96,
-    temperature: float = 0.7,
-    min_p: float = 0.05,
-    no_repeat_ngram_size: int = 3,
-    repetition_penalty: float = 1.2,
-    use_anchor: bool = True,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    repetition_penalty: float,
+    no_repeat_ngram_size: int,
+    use_anchor: bool,
 ):
-    """
-    gr.ChatInterface-compatible streaming generator. Yields incrementally
-    as tokens are produced, and appends a real (measured, not decorative)
-    telemetry line -- token count / time-to-first-token / tokens-per-
-    second -- under the reply, matching the "signature element" of this
-    UI: honest, live numbers about what the model is actually doing.
-
-    IMPORTANT, honest by design: `history` is accepted (gr.ChatInterface
-    requires the signature) and IS shown to the person as an ongoing
-    conversation, but is deliberately NOT concatenated into the prompt
-    fed to the model. Arya is a continuation-only base model, never
-    instruction-tuned on multi-turn User/Assistant-style dialogue at
-    scale -- feeding it a growing instruction-formatted transcript would
-    push it further out of its training distribution with every turn,
-    not closer to coherence. Each message is generated as its own fresh
-    continuation (with the anchor primer). This is stated plainly in the
-    UI caption rather than left as a surprise.
-    """
-    try:
-        full_prompt = (ANCHOR_TEXT + message) if use_anchor else message
+    if not message.strip():
+        yield "Please enter a prompt."
+        return
+
+    prompt      = (_ANCHOR + message) if use_anchor else message
+    prompt_ids  = _tokenizer.encode(prompt, add_bos=True, add_eos=False)
+    if len(prompt_ids) >= _ctx_len:
+        prompt_ids = prompt_ids[-(_ctx_len - 1):]
+
+    generated_ids : list[int] = []
+    eos_id        = getattr(_tokenizer, "eos_id", None)
+    t_start       = time.perf_counter()
+    ttft          : float | None = None
+
+    def _stats(n: int) -> str:
+        elapsed = max(time.perf_counter() - t_start, 1e-6)
+        ttft_ms = (ttft * 1000) if ttft is not None else 0.0
+        return (
+            f"\n\n<span style='font-size:11px;font-family:\"IBM Plex Mono\",monospace;"
+            f"color:#9ca3af;letter-spacing:0.03em'>"
+            f"{n} tokens · TTFT {ttft_ms:.0f} ms · {n/elapsed:.1f} tok/s</span>"
+        )
 
-        tokens = tokenizer.encode(full_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:]
+    idx = torch.tensor([prompt_ids], dtype=torch.long, device=_device)
 
-        generated_ids: list = []
-        t_start = time.perf_counter()
-        ttft = None
+    with _gen_lock, torch.no_grad(), _cast_ctx():
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= _ctx_len else idx[:, -_ctx_len:]
+            logits, _ = _model(idx_cond)
+            logits     = logits[:, -1, :].float()
 
-        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,
-        )
+            row = logits[0:1]
+            row = _apply_repetition_penalty(row, generated_ids, repetition_penalty)
+            row = _block_repeated_ngrams(row, generated_ids, no_repeat_ngram_size)
+            row = row / max(temperature, 1e-5)
+            if top_k > 0:
+                row = _top_k_filter(row, top_k)
+            if top_p < 1.0:
+                row = _top_p_filter(row, top_p)
+
+            probs   = torch.softmax(row, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1).item()
+
+            if ttft is None:
+                ttft = time.perf_counter() - t_start
+
+            generated_ids.append(next_id)
+            idx = torch.cat([idx, torch.tensor([[next_id]], device=_device)], dim=1)
+
+            if eos_id is not None and next_id == eos_id:
+                break
+
+            full = _tokenizer.decode(generated_ids)
+            yield full + _stats(len(generated_ids))
+
+    final = _tokenizer.decode(generated_ids) or "(no output)"
+    yield final + _stats(len(generated_ids))
 
-        logits_row = np.array(logits[0, -1, :], dtype=np.float32, copy=True)
-        logits_row = apply_repetition_controls(
-            logits_row, generated_ids, no_repeat_ngram_size, repetition_penalty
-        )
-        next_token = sample_from_logits(logits_row, temperature=temperature, min_p=min_p)
-        tokens.append(next_token)
-        generated_ids.append(next_token)
-        ttft = time.perf_counter() - t_start
-
-        partial_text = tokenizer.decode(generated_ids)
-        yield _with_telemetry(partial_text, len(generated_ids), ttft, t_start)
-
-        if next_token != tokenizer.eos_id and len(tokens) < CONTEXT:
-            cache_pos = len(prompt_ids[0])
-
-            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,
-                )
-                cache_pos += 1
-
-                logits_row = np.array(logits[0, -1, :], dtype=np.float32, copy=True)
-                logits_row = apply_repetition_controls(
-                    logits_row, generated_ids, no_repeat_ngram_size, repetition_penalty
-                )
-                next_token = sample_from_logits(logits_row, temperature=temperature, min_p=min_p)
-                tokens.append(next_token)
-                generated_ids.append(next_token)
-
-                if next_token == tokenizer.eos_id:
-                    break
-
-                partial_text = tokenizer.decode(generated_ids)
-                yield _with_telemetry(partial_text, len(generated_ids), ttft, t_start)
-
-                if len(tokens) >= CONTEXT:
-                    break
-
-    except Exception as e:
-        yield f"Error: {str(e)}"
-
-
-def _with_telemetry(text: str, n_tok: int, ttft: float, t_start: float) -> str:
-    elapsed = max(time.perf_counter() - t_start, 1e-6)
-    tok_per_s = n_tok / elapsed
-    telemetry = (
-        f"\n\n<sub style='font-family:monospace;color:var(--body-text-color-subdued)'>"
-        f"{n_tok} tok · TTFT {ttft * 1000:.0f} ms · {tok_per_s:.1f} tok/s</sub>"
-    )
-    return text + telemetry
 
+# ── UI ────────────────────────────────────────────────────────────────────────
 
-# ============================================================
-# Gradio UI — Arya
-# ============================================================
-# NOTE on "gr.server()": there is no such function in Gradio. gr.Blocks
-# and gr.ChatInterface are ALREADY built on FastAPI/Uvicorn internally --
-# demo.launch() below is what actually starts that server. Same
-# clarification as the earlier "use FastAPI for speed" question: the web
-# framework was never the bottleneck for a custom JAX model like this.
+CSS = """
+/* ══════════════════════════════════════════════════════
+   Sarvam-style: warm white · near-black · saffron accent
+   Inter (body) + IBM Plex Mono (code / telemetry)
+   ══════════════════════════════════════════════════════ */
+
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500&display=swap');
 
-ARYA_CSS = """
 :root {
-    --arya-bg: #100C1A;
-    --arya-bg-2: #17111F;
-    --arya-panel: #1D1730;
-    --arya-panel-2: #271F42;
-    --arya-accent: #E8A33D;
-    --arya-accent-soft: #F0BE73;
-    --arya-accent-2: #5FBFA8;
-    --arya-indigo: #5B67D9;
-    --arya-text: #F3EEE3;
-    --arya-text-dim: #948BAA;
-    --arya-hairline: #2A2440;
-
-    /* Retheme Gradio's OWN documented CSS variables (verified against the
-       installed theme's actual variable names) instead of guessing at
-       hashed/scoped internal component classes -- this cascades correctly
-       into sliders, checkboxes, accordions, buttons, etc. automatically,
-       and won't silently stop working on a Gradio version bump the way
-       hand-guessed internal class names would. */
-    --body-background-fill: var(--arya-bg);
-    --background-fill-primary: var(--arya-bg);
-    --background-fill-secondary: var(--arya-panel);
-    --block-background-fill: var(--arya-panel);
-    --block-border-color: var(--arya-hairline);
-    --block-title-text-color: var(--arya-text-dim);
-    --block-label-text-color: var(--arya-text-dim);
-    --body-text-color: var(--arya-text);
-    --body-text-color-subdued: var(--arya-text-dim);
-    --border-color-primary: var(--arya-hairline);
-    --border-color-accent: var(--arya-accent);
-    --color-accent: var(--arya-accent);
-    --color-accent-soft: color-mix(in srgb, var(--arya-accent) 20%, transparent);
-    --input-background-fill: var(--arya-panel-2);
-    --input-border-color: var(--arya-hairline);
-    --input-border-color-focus: var(--arya-accent);
-    --input-placeholder-color: var(--arya-text-dim);
-    --slider-color: var(--arya-accent);
-    --checkbox-background-color-selected: var(--arya-accent);
-    --checkbox-border-color-selected: var(--arya-accent);
-    --checkbox-label-background-fill-selected: color-mix(in srgb, var(--arya-accent) 18%, var(--arya-panel));
-    --button-primary-background-fill: linear-gradient(135deg, var(--arya-accent), var(--arya-accent-soft));
-    --button-primary-background-fill-hover: var(--arya-accent-soft);
-    --button-primary-text-color: #201306;
-    --button-secondary-background-fill: var(--arya-panel-2);
-    --button-secondary-border-color: var(--arya-hairline);
-    --button-secondary-background-fill-hover: var(--arya-panel);
-    --link-text-color: var(--arya-accent-2);
-}
-
-* { scrollbar-width: thin; scrollbar-color: var(--arya-panel-2) transparent; }
-::-webkit-scrollbar { width: 8px; height: 8px; }
-::-webkit-scrollbar-thumb { background: var(--arya-panel-2); border-radius: 8px; }
-::-webkit-scrollbar-thumb:hover { background: var(--arya-hairline); }
+    /* surface */
+    --s-bg:        #fafaf9;
+    --s-white:     #ffffff;
+    --s-warm-50:   #f5f5f4;
+    --s-warm-100:  #e7e5e4;
+    --s-warm-200:  #d6d3d1;
+
+    /* text */
+    --s-ink:       #0a0a0a;
+    --s-ink-2:     #292524;
+    --s-muted:     #78716c;
+    --s-faint:     #a8a29e;
+
+    /* accent — saffron */
+    --s-saf:       #ea580c;
+    --s-saf-hover: #c2410c;
+    --s-saf-light: #fff7ed;
+    --s-saf-mid:   #fed7aa;
+
+    /* Gradio token overrides */
+    --body-background-fill:          var(--s-bg);
+    --background-fill-primary:       var(--s-white);
+    --background-fill-secondary:     var(--s-warm-50);
+    --block-background-fill:         var(--s-white);
+    --block-border-color:            var(--s-warm-100);
+    --block-title-text-color:        var(--s-muted);
+    --block-label-text-color:        var(--s-muted);
+    --body-text-color:               var(--s-ink);
+    --body-text-color-subdued:       var(--s-muted);
+    --border-color-primary:          var(--s-warm-100);
+    --border-color-accent:           var(--s-saf);
+    --color-accent:                  var(--s-saf);
+    --input-background-fill:         var(--s-white);
+    --input-border-color:            var(--s-warm-200);
+    --input-border-color-focus:      var(--s-saf);
+    --input-placeholder-color:       var(--s-faint);
+    --slider-color:                  var(--s-saf);
+    --checkbox-background-color-selected:   var(--s-saf);
+    --checkbox-border-color-selected:       var(--s-saf);
+    --checkbox-label-background-fill-selected:
+        color-mix(in srgb, var(--s-saf) 8%, var(--s-white));
+    --button-primary-background-fill:       var(--s-saf);
+    --button-primary-background-fill-hover: var(--s-saf-hover);
+    --button-primary-text-color:            #ffffff;
+    --button-secondary-background-fill:     var(--s-white);
+    --button-secondary-border-color:        var(--s-warm-200);
+    --button-secondary-background-fill-hover: var(--s-warm-50);
+    --link-text-color:                      var(--s-saf);
+}
+
+/* ── Base ──────────────────────────────────────────────── */
+* { box-sizing: border-box; }
+
+body, .gradio-container {
+    font-family: 'Inter', system-ui, -apple-system, sans-serif !important;
+    background: var(--s-bg) !important;
+    color: var(--s-ink) !important;
+}
 
 .gradio-container {
-    background:
-        radial-gradient(ellipse 900px 400px at 15% -10%, color-mix(in srgb, var(--arya-accent) 10%, transparent), transparent 60%),
-        radial-gradient(ellipse 700px 400px at 100% 0%, color-mix(in srgb, var(--arya-indigo) 12%, transparent), transparent 55%),
-        var(--arya-bg) !important;
-    font-family: 'Inter', system-ui, sans-serif !important;
-    max-width: 880px !important;
+    max-width: 820px !important;
+    margin: 0 auto !important;
+    padding: 0 20px !important;
 }
 
-/* ── Topbar ── */
-#arya-topbar {
+/* ── Nav bar ──────────────────────────────────────────── */
+#s-nav {
     display: flex;
     align-items: center;
+    padding: 18px 0 16px;
+    border-bottom: 1px solid var(--s-warm-100);
+    margin-bottom: 0;
     gap: 12px;
-    padding: 18px 4px 16px 4px;
-    border-bottom: 1px solid var(--arya-hairline);
-    margin-bottom: 10px;
-}
-#arya-mark {
-    width: 34px; height: 34px; flex-shrink: 0;
-    display: flex; align-items: center; justify-content: center;
-    border-radius: 10px;
-    background: linear-gradient(135deg, var(--arya-accent), var(--arya-indigo));
-    box-shadow: 0 2px 14px color-mix(in srgb, var(--arya-accent) 35%, transparent);
-}
-#arya-mark svg { width: 19px; height: 19px; }
-#arya-title-block { display: flex; flex-direction: column; line-height: 1.15; }
-#arya-title {
-    font-family: 'IBM Plex Mono', monospace;
+}
+#s-wordmark {
+    font-family: 'Inter', sans-serif;
     font-weight: 700;
-    font-size: 1.18rem;
-    letter-spacing: 0.14em;
-    color: var(--arya-text);
+    font-size: 1.05rem;
+    letter-spacing: -0.02em;
+    color: var(--s-ink);
+    display: flex;
+    align-items: center;
+    gap: 8px;
+    text-decoration: none;
 }
-#arya-subtitle {
-    font-family: 'IBM Plex Mono', monospace;
-    font-size: 0.66rem;
-    letter-spacing: 0.06em;
-    color: var(--arya-text-dim);
+#s-wordmark-dot {
+    width: 8px; height: 8px; border-radius: 50%;
+    background: var(--s-saf);
+    flex-shrink: 0;
 }
-#arya-status-chip {
-    font-family: 'IBM Plex Mono', monospace;
+#s-nav-tag {
     font-size: 0.72rem;
     font-weight: 500;
-    color: var(--arya-accent-2);
-    display: flex; align-items: center; gap: 6px;
-    padding: 5px 10px;
+    color: var(--s-saf);
+    background: var(--s-saf-light);
+    border: 1px solid var(--s-saf-mid);
+    padding: 3px 9px;
     border-radius: 999px;
-    background: color-mix(in srgb, var(--arya-accent-2) 12%, transparent);
-    border: 1px solid color-mix(in srgb, var(--arya-accent-2) 30%, transparent);
-    margin-left: 14px;
-}
-#arya-status-chip::before {
-    content: "";
-    width: 6px; height: 6px; border-radius: 50%;
-    background: var(--arya-accent-2);
-    box-shadow: 0 0 8px var(--arya-accent-2);
-    animation: arya-pulse 2s ease-in-out infinite;
-}
-@keyframes arya-pulse {
-    0%, 100% { opacity: 1; transform: scale(1); }
-    50% { opacity: 0.5; transform: scale(0.8); }
-}
-#arya-specs {
+    letter-spacing: 0.01em;
+}
+#s-nav-right {
     margin-left: auto;
-    display: flex; gap: 6px; flex-wrap: wrap; justify-content: flex-end;
+    display: flex;
+    align-items: center;
+    gap: 20px;
 }
-.arya-chip {
-    font-family: 'IBM Plex Mono', monospace;
-    font-size: 0.66rem;
-    letter-spacing: 0.03em;
-    color: var(--arya-text-dim);
-    padding: 4px 9px;
-    border-radius: 6px;
-    background: var(--arya-panel);
-    border: 1px solid var(--arya-hairline);
-    white-space: nowrap;
+.s-nav-link {
+    font-size: 0.82rem;
+    color: var(--s-muted);
+    text-decoration: none;
+    font-weight: 500;
+    transition: color 0.12s;
 }
+.s-nav-link:hover { color: var(--s-ink); }
 
-/* ── Chat transcript (verified real Gradio 6 Chatbot classes, not guessed:
-     .user-row / .bot-row / .message-bubble-border confirmed present in the
-     installed package's compiled component CSS) ── */
+/* ── Hero ─────────────────────────────────────────────── */
+#s-hero {
+    padding: 56px 0 40px;
+    border-bottom: 1px solid var(--s-warm-100);
+}
+#s-hero-eyebrow {
+    font-size: 0.78rem;
+    font-weight: 600;
+    letter-spacing: 0.06em;
+    text-transform: uppercase;
+    color: var(--s-saf);
+    margin-bottom: 16px;
+}
+#s-hero h1 {
+    font-size: clamp(2rem, 5vw, 2.8rem);
+    font-weight: 700;
+    letter-spacing: -0.03em;
+    line-height: 1.12;
+    color: var(--s-ink);
+    margin: 0 0 16px;
+    max-width: 560px;
+}
+#s-hero-desc {
+    font-size: 1rem;
+    color: var(--s-muted);
+    line-height: 1.65;
+    max-width: 480px;
+    margin-bottom: 28px;
+    font-weight: 400;
+}
+#s-stat-row {
+    display: flex;
+    gap: 32px;
+    flex-wrap: wrap;
+}
+.s-stat {
+    display: flex;
+    flex-direction: column;
+    gap: 2px;
+}
+.s-stat-val {
+    font-size: 1.1rem;
+    font-weight: 700;
+    letter-spacing: -0.02em;
+    color: var(--s-ink);
+    font-variant-numeric: tabular-nums;
+}
+.s-stat-lbl {
+    font-size: 0.72rem;
+    color: var(--s-faint);
+    font-weight: 500;
+    text-transform: uppercase;
+    letter-spacing: 0.05em;
+}
+
+/* ── Section label ────────────────────────────────────── */
+.s-section-label {
+    font-size: 0.72rem;
+    font-weight: 600;
+    letter-spacing: 0.08em;
+    text-transform: uppercase;
+    color: var(--s-faint);
+    padding: 28px 0 12px;
+    border-bottom: none;
+}
+
+/* ── Chat bubbles ─────────────────────────────────────── */
+/* User: saffron-tinted */
 .user-row {
-    background: linear-gradient(135deg, var(--arya-accent), var(--arya-accent-soft)) !important;
-    color: #201306 !important;
+    background: var(--s-saf) !important;
+    border-radius: 12px 12px 2px 12px !important;
 }
-.user-row * { color: #201306 !important; }
+.user-row * { color: #fff !important; }
+
+/* Bot: white card with warm border */
 .bot-row {
-    background: var(--arya-panel) !important;
-    border: 1px solid var(--arya-hairline) !important;
+    background: var(--s-white) !important;
+    border: 1px solid var(--s-warm-100) !important;
+    border-radius: 12px 12px 12px 2px !important;
 }
+
 .message-bubble-border {
-    border-color: var(--arya-hairline) !important;
-    box-shadow: 0 4px 18px rgba(0,0,0,0.25);
-    animation: arya-rise 0.25s ease-out;
-}
-@keyframes arya-rise {
-    from { opacity: 0; transform: translateY(6px); }
-    to { opacity: 1; transform: translateY(0); }
+    box-shadow: 0 1px 3px rgba(0,0,0,0.06) !important;
 }
+
+/* Code in bot replies */
 .bot-row code, .bot-row pre {
     font-family: 'IBM Plex Mono', monospace !important;
-    background: var(--arya-bg-2) !important;
-    border: 1px solid var(--arya-hairline) !important;
+    background: var(--s-warm-50) !important;
+    border: 1px solid var(--s-warm-100) !important;
     border-radius: 6px !important;
+    font-size: 0.82rem !important;
 }
 
-/* ── Input row ── */
-textarea, input[type="text"] {
-    transition: border-color 0.15s ease, box-shadow 0.15s ease !important;
+/* ── Textbox & send button ────────────────────────────── */
+textarea {
+    font-family: 'Inter', sans-serif !important;
+    font-size: 0.95rem !important;
+    border-radius: 8px !important;
+    border: 1.5px solid var(--s-warm-200) !important;
+    transition: border-color 0.15s, box-shadow 0.15s !important;
+    resize: none !important;
+    background: var(--s-white) !important;
+    color: var(--s-ink) !important;
 }
-textarea:focus, input[type="text"]:focus {
-    box-shadow: 0 0 0 3px color-mix(in srgb, var(--arya-accent) 22%, transparent) !important;
+textarea:focus {
+    border-color: var(--s-saf) !important;
+    box-shadow: 0 0 0 3px color-mix(in srgb, var(--s-saf) 12%, transparent) !important;
+    outline: none !important;
 }
+textarea::placeholder { color: var(--s-faint) !important; }
+
 button.primary {
-    transition: transform 0.12s ease, box-shadow 0.12s ease !important;
-    box-shadow: 0 2px 12px color-mix(in srgb, var(--arya-accent) 30%, transparent) !important;
+    background: var(--s-saf) !important;
+    border: none !important;
+    border-radius: 8px !important;
+    font-weight: 600 !important;
+    font-size: 0.88rem !important;
+    letter-spacing: 0.01em !important;
+    color: #fff !important;
+    transition: background 0.12s, transform 0.1s !important;
+}
+button.primary:hover {
+    background: var(--s-saf-hover) !important;
+    transform: translateY(-1px) !important;
 }
-button.primary:hover { transform: translateY(-1px) scale(1.03); }
 
-/* ── Settings accordion & example chips: themed via the CSS variables
-   above, no internal class names needed ── */
+/* ── Settings accordion ───────────────────────────────── */
+.accordion {
+    border: 1px solid var(--s-warm-100) !important;
+    border-radius: 10px !important;
+    background: var(--s-white) !important;
+    overflow: hidden;
+}
+.accordion > .label-wrap {
+    padding: 12px 16px !important;
+    font-size: 0.85rem !important;
+    font-weight: 600 !important;
+    color: var(--s-ink-2) !important;
+    background: var(--s-warm-50) !important;
+    border-bottom: 1px solid var(--s-warm-100) !important;
+}
+
+/* Sliders */
+input[type="range"] { accent-color: var(--s-saf) !important; }
+
+/* ── Example chips ────────────────────────────────────── */
 .example {
-    transition: transform 0.12s ease, border-color 0.12s ease !important;
+    font-size: 0.82rem !important;
+    font-weight: 500 !important;
+    border: 1px solid var(--s-warm-200) !important;
+    border-radius: 6px !important;
+    color: var(--s-ink-2) !important;
+    background: var(--s-white) !important;
+    padding: 6px 12px !important;
+    transition: border-color 0.12s, color 0.12s, background 0.12s !important;
 }
 .example:hover {
-    transform: translateY(-2px);
-    border-color: var(--arya-accent) !important;
+    border-color: var(--s-saf) !important;
+    color: var(--s-saf) !important;
+    background: var(--s-saf-light) !important;
 }
 
-#arya-caption {
-    font-family: 'IBM Plex Mono', monospace;
+/* ── Footer ───────────────────────────────────────────── */
+#s-footer {
+    border-top: 1px solid var(--s-warm-100);
+    padding: 28px 0 20px;
+    display: flex;
+    justify-content: space-between;
+    align-items: flex-start;
+    flex-wrap: wrap;
+    gap: 16px;
+    margin-top: 16px;
+}
+#s-footer-left {
+    display: flex;
+    flex-direction: column;
+    gap: 5px;
+}
+#s-footer-brand {
+    font-weight: 700;
+    font-size: 0.88rem;
+    color: var(--s-ink);
+    letter-spacing: -0.01em;
+}
+#s-footer-copy {
+    font-size: 0.72rem;
+    color: var(--s-faint);
+}
+#s-footer-chips {
+    display: flex;
+    gap: 6px;
+    flex-wrap: wrap;
+    align-items: center;
+}
+.s-badge {
     font-size: 0.68rem;
-    color: var(--arya-text-dim);
-    text-align: center;
-    padding: 16px 20px 6px 20px;
-    letter-spacing: 0.03em;
-    border-top: 1px solid var(--arya-hairline);
-    margin-top: 12px;
+    font-weight: 500;
+    font-family: 'IBM Plex Mono', monospace;
+    color: var(--s-muted);
+    background: var(--s-warm-50);
+    border: 1px solid var(--s-warm-100);
+    padding: 3px 8px;
+    border-radius: 4px;
+    letter-spacing: 0.02em;
+    white-space: nowrap;
+}
+.s-badge-accent {
+    color: var(--s-saf);
+    background: var(--s-saf-light);
+    border-color: var(--s-saf-mid);
 }
-"""
-
-ARYA_HEAD = """
-<link rel="preconnect" href="https://fonts.googleapis.com">
-<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500;700&display=swap" rel="stylesheet">
-"""
 
-# Simple geometric monogram, not a decorative flourish -- the gradient fill
-# (accent -> indigo) matches #arya-mark's CSS background, single mark
-# rather than a busier logo, keeps it legible at 19px.
-ARYA_MARK_SVG = """
-<svg viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
-  <path d="M12 3L21 20H16.5L12 11L7.5 20H3L12 3Z" fill="#100C1A"/>
-</svg>
+/* ── Misc fixes ───────────────────────────────────────── */
+* { scrollbar-width: thin; scrollbar-color: var(--s-warm-200) transparent; }
+::-webkit-scrollbar { width: 5px; }
+::-webkit-scrollbar-thumb { background: var(--s-warm-200); border-radius: 4px; }
 """
 
 
 def main():
-    spec_chips = [
-        f"{model.count_params() / 1e6:.1f}M PARAMS",
-        f"{numberoflayers}L",
-        f"CTX {CONTEXT}",
-        f"GQA {numberofheads}/{num_kv_heads}",
-    ]
-    spec_html = "".join(f'<span class="arya-chip">{c}</span>' for c in spec_chips)
-
-    with gr.Blocks(title="Arya") as demo:
-
-        gr.HTML(f"""
-        <div id="arya-topbar">
-            <div id="arya-mark">{ARYA_MARK_SVG}</div>
-            <div id="arya-title-block">
-                <span id="arya-title">ARYA</span>
-                <span id="arya-subtitle">FROM-SCRATCH SMALL LM</span>
+    # Build runtime stats for the hero section
+    device_label  = str(_device).upper()
+    dtype_label   = str(_cast_dtype).replace("torch.", "").upper() if _cast_dtype else "FP32"
+    layer_str     = str(_n_layer)
+    head_str      = f"{_n_head}/{_kv_heads}"
+
+    header_html = f"""
+    <div id="s-nav">
+        <span id="s-wordmark">
+            <span id="s-wordmark-dot"></span>
+            Pragya
+        </span>
+        <span id="s-nav-tag">by Arush Kumar</span>
+        <div id="s-nav-right">
+            <span class="s-nav-link">PyTorch</span>
+            <span class="s-nav-link">{device_label}</span>
+        </div>
+    </div>
+
+    <div id="s-hero">
+        <div id="s-hero-eyebrow">Arush Kumar · From-scratch language model</div>
+        <h1>Pragya — a model<br>built from the ground up.</h1>
+        <p id="s-hero-desc">
+            A custom GPT-style architecture with Mixture-of-Depths routing,
+            Multi-Token Prediction heads, and YaRN RoPE context scaling.
+            Every weight trained from scratch in PyTorch. Type a prompt and let it continue.
+        </p>
+        <div id="s-stat-row">
+            <div class="s-stat">
+                <span class="s-stat-val">{_n_params_str}</span>
+                <span class="s-stat-lbl">Parameters</span>
+            </div>
+            <div class="s-stat">
+                <span class="s-stat-val">{_ctx_len:,}</span>
+                <span class="s-stat-lbl">Context length</span>
+            </div>
+            <div class="s-stat">
+                <span class="s-stat-val">{layer_str}L · {head_str}H</span>
+                <span class="s-stat-lbl">Depth · GQA heads</span>
+            </div>
+            <div class="s-stat">
+                <span class="s-stat-val">{_pattern}</span>
+                <span class="s-stat-lbl">Attention pattern</span>
+            </div>
+            <div class="s-stat">
+                <span class="s-stat-val">{dtype_label}</span>
+                <span class="s-stat-lbl">Inference dtype</span>
             </div>
-            <span id="arya-status-chip">READY</span>
-            <div id="arya-specs">{spec_html}</div>
         </div>
-        """)
+    </div>
+    """
+
+    footer_html = f"""
+    <div id="s-footer">
+        <div id="s-footer-left">
+            <span id="s-footer-brand">Pragya &nbsp;·&nbsp; <span style="font-weight:400;color:var(--s-muted)">Arush Kumar</span></span>
+            <span id="s-footer-copy">
+                Continuation model · conversation history is shown in UI but not fed back to the model ·
+                TTFT and tok/s are live-measured each generation
+            </span>
+        </div>
+        <div id="s-footer-chips">
+            <span class="s-badge s-badge-accent">{_n_params_str} params</span>
+            <span class="s-badge">{device_label}</span>
+            <span class="s-badge">{dtype_label}</span>
+            <span class="s-badge">ctx {_ctx_len}</span>
+            <span class="s-badge">{_pattern}</span>
+        </div>
+    </div>
+    """
+
+    with gr.Blocks(title="Pragya by Arush Kumar") as demo:
+        gr.HTML(header_html)
+
+        gr.HTML('<div class="s-section-label">Try it</div>')
 
         gr.ChatInterface(
             fn=respond,
             additional_inputs=[
-                gr.Slider(10, 256, value=96, step=10, label="Max tokens"),
-                gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
-                gr.Slider(0.0, 0.3, value=0.05, step=0.01, label="Min-P",
-                          info="Replaces top-k. Higher = stricter."),
-                gr.Slider(0, 6, value=3, step=1, label="No-repeat n-gram size",
-                          info="Hard-blocks repeating any phrase this long. 0 disables."),
-                gr.Slider(1.0, 1.5, value=1.2, step=0.05, label="Repetition penalty"),
+                gr.Slider(16, _MAX_TOKENS_CAP, value=_DEFAULT_TOKENS, step=8,
+                          label="Max new tokens"),
+                gr.Slider(0.1, 2.0, value=0.8, step=0.05,
+                          label="Temperature"),
+                gr.Slider(0, 200, value=50, step=5,
+                          label="Top-k",
+                          info="0 disables top-k filtering."),
+                gr.Slider(0.5, 1.0, value=1.0, step=0.01,
+                          label="Top-p (nucleus)",
+                          info="1.0 disables nucleus sampling."),
+                gr.Slider(1.0, 2.0, value=1.3, step=0.05,
+                          label="Repetition penalty",
+                          info="Count-scaled: penalty^occurrences. 1.0 = off."),
+                gr.Slider(0, 6, value=3, step=1,
+                          label="No-repeat n-gram",
+                          info="Hard-bans any n-gram of this length from repeating. 0 = off."),
                 gr.Checkbox(value=True, label="Narrative anchor",
-                            info="Primes the model with clean prose before your prompt."),
+                            info="Prepends a short coherent prose passage to prime the context."),
             ],
-            additional_inputs_accordion=gr.Accordion("Settings", open=False),
+            additional_inputs_accordion=gr.Accordion("Generation settings", open=False),
+            chatbot=gr.Chatbot(
+                height=440,
+                show_label=False,
+                render_markdown=True,
+                bubble_full_width=False,
+            ),
+            textbox=gr.Textbox(
+                placeholder="Start a sentence and let Pragya complete it …",
+                show_label=False,
+                lines=2,
+            ),
             examples=[
-                ["Once upon a time, in a village by the mountains,"],
-                ["The scientist looked at the data and realized"],
-                ["Explain photosynthesis simply."],
+                ["The scientist stared at the anomalous readings and realized"],
+                ["Once upon a time, in a village at the foot of a mountain,"],
+                ["The rules of the game were simple, but nobody had read them."],
+                ["Deep in the forest, where the trees grew so tall they blocked the sky,"],
             ],
-            chatbot=gr.Chatbot(height=460, show_label=False),
-            textbox=gr.Textbox(placeholder="Ask Arya anything...", show_label=False),
+            cache_examples=False,
         )
 
-        gr.HTML("""
-        <div id="arya-caption">
-            CONTINUATION MODEL, NOT INSTRUCTION-TUNED — EACH REPLY IS A FRESH
-            GENERATION, NOT A REMEMBERED CONVERSATION · RUNS FULLY ON YOUR HARDWARE
-        </div>
-        """)
+        gr.HTML(footer_html)
 
-    # css/head/theme live on launch(), not Blocks(), as of Gradio 6 --
-    # verified against the actually-installed version rather than assumed.
     demo.launch(
-        share=True, server_name="0.0.0.0", server_port=7860,
-        css=ARYA_CSS, head=ARYA_HEAD,
-        theme=gr.themes.Base(primary_hue=gr.themes.colors.orange,
-                              neutral_hue=gr.themes.colors.slate),
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        css=CSS,
+        theme=gr.themes.Base(
+            primary_hue=gr.themes.colors.orange,
+            neutral_hue=gr.themes.colors.stone,
+            font=gr.themes.GoogleFont("Inter"),
+        ),
     )
 
+
 if __name__ == "__main__":
     main()