Update app.py

app.py

@@ -13,16 +13,24 @@ Environment variables (Space Settings -> Variables):
     TOKENIZER_PATH       tokenizer.pkl or its folder     (default: auto-detected next to app.py)
     MAX_TOKENS_CAP       slider maximum                   (default 512)
     DEFAULT_MAX_TOKENS   slider default                   (default 128)
+
+UI: custom Blocks chat -- send button, interruptible generation (Stop), quick-start
+chips, and a live stats line (TTFT / tok/s) under the conversation. Works on Gradio
+4.x and 5.x; uses the messages chat format when available, falls back to tuples.
 """
 from __future__ import annotations
 
 import contextlib
+import html as _htmllib
+import inspect
 import os
 import queue
 import sys
+import tempfile
 import threading
 import time
 from pathlib import Path
+from typing import Any
 
 import gradio as gr
 import torch
@@ -31,14 +39,12 @@ _HERE = Path(__file__).resolve().parent
 if str(_HERE) not in sys.path:
     sys.path.insert(0, str(_HERE))
 
-import inspect  # noqa: E402
-
 from inference import generate, load_model, pick_device, _resolve_dtype  # noqa: E402
 
 if "on_text" not in inspect.signature(generate).parameters:
     raise RuntimeError("inference.py on this Space is out of date (generate() has no on_text argument). "
                        "Upload the current inference.py next to app.py.")
-from tokenizer import TokenizerWrapper                                   # noqa: E402
+from tokenizer import TokenizerWrapper  # noqa: E402
 
 _CHECKPOINT_LOCAL = os.environ.get("CHECKPOINT_PATH", str(_HERE / "pragya.pt"))
 _HF_REPO          = os.environ.get("CHECKPOINT_REPO", "")
@@ -93,10 +99,33 @@ _dtype     = _resolve_dtype("auto", _device)          # autocast dtype, None = f
 _dtype_lbl = str(_dtype).replace("torch.", "").upper() if _dtype else "FP32"
 _device_lbl = str(_device).upper()
 _gen_lock  = threading.Lock()        # one generation at a time on the shared model
+_ui_guard  = threading.Lock()        # one respond() stream at a time in the UI
 
 print(f"[App] Ready: {_n_params} params, ctx={_ctx_len}, device={_device}", flush=True)
 
 
+# -- avatars (cosmetic; PIL ships with gradio) --------------------------------
+def _make_avatar(bg: tuple[int, int, int], dot: tuple[int, int, int] | None) -> str:
+    from PIL import Image, ImageDraw
+    size, r = 256, 46
+    img = Image.new("RGBA", (size, size), (0, 0, 0, 0))
+    draw = ImageDraw.Draw(img)
+    draw.ellipse((6, 6, size - 6, size - 6), fill=(*bg, 255))
+    if dot:
+        draw.ellipse((size / 2 - r, size / 2 - r, size / 2 + r, size / 2 + r), fill=(*dot, 255))
+    path = Path(tempfile.gettempdir()) / f"pragya_avatar_{bg[0]}_{bg[1]}_{bg[2]}_{bool(dot)}.png"
+    img.save(path)
+    return str(path)
+
+
+try:
+    _AVATARS = (_make_avatar((41, 37, 36), (168, 162, 158)),    # user
+                _make_avatar((234, 88, 12), (255, 255, 255)))   # Pragya
+except Exception as _avatar_exc:                                # never block startup over avatars
+    print(f"[App] Avatars skipped ({_avatar_exc}).", flush=True)
+    _AVATARS = None
+
+
 # -- generation ---------------------------------------------------------------
 _ANCHOR = (
     "The old lighthouse stood at the edge of the cliff, its beam sweeping slowly across the dark "
@@ -105,69 +134,218 @@ _ANCHOR = (
 )
 
 
-def _stats_html(n_tokens: int, ttft_ms: float | None, elapsed: float) -> str:
-    rate = n_tokens / max(elapsed, 1e-6)
-    ttft = f"{ttft_ms:.0f} ms" if ttft_ms is not None else "-"
-    return (
-        "<div class='p-stats'>"
-        f"{n_tokens} tokens &middot; TTFT {ttft} &middot; {rate:.1f} tok/s"
-        "</div>"
-    )
+class _GenerationStopped(Exception):
+    """Raised inside the on_text callback to abort generation mid-stream."""
+
+
+_stop_lock = threading.Lock()
+_active_stop: threading.Event | None = None
+
+
+def _new_stop_event() -> threading.Event:
+    global _active_stop
+    ev = threading.Event()
+    with _stop_lock:
+        _active_stop = ev
+    return ev
 
 
+def _request_stop() -> None:
+    with _stop_lock:
+        ev = _active_stop
+    if ev is not None:
+        ev.set()
+
+
+def _count(text: str) -> int:
+    return len(_tokenizer.encode(text, add_bos=False, add_eos=False)) if text else 0
+
+
+# -- chat-format helpers (messages format when supported, tuples otherwise) ---
+_MSG_FORMAT: bool = True
+
+
+def _text(content: Any) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, (list, tuple)):
+        parts = []
+        for part in content:
+            parts.append(str(part.get("text", "")) if isinstance(part, dict) else str(part))
+        return "".join(parts)
+    return str(content)
+
+
+def _norm_history(value: Any) -> list[dict[str, str]]:
+    """Whatever the Chatbot hands us (message dicts or (user, bot) tuples) -> message dicts."""
+    messages: list[dict[str, str]] = []
+    for item in value or []:
+        if isinstance(item, dict):
+            messages.append({"role": str(item.get("role", "assistant")),
+                             "content": _text(item.get("content"))})
+        elif isinstance(item, (list, tuple)) and len(item) == 2:
+            user, bot = item
+            if user is not None:
+                messages.append({"role": "user", "content": _text(user)})
+            if bot is not None:
+                messages.append({"role": "assistant", "content": _text(bot)})
+    return messages
+
+
+def _chat_value(messages: list[dict[str, str]]) -> Any:
+    """Internal message dicts -> a value the Chatbot component understands."""
+    if _MSG_FORMAT:
+        return messages
+    pairs: list[list[Any]] = []
+    pending_user: str | None = None
+    for message in messages:
+        if message["role"] == "user":
+            pending_user = message["content"]
+        else:
+            pairs.append([pending_user or "", message["content"]])
+            pending_user = None
+    if pending_user is not None:
+        pairs.append([pending_user, None])
+    return pairs
+
+
+# -- status line ---------------------------------------------------------------
+def _status_html(kind: str, n_tokens: int = 0, ttft_ms: float | None = None,
+                 elapsed: float = 0.0, note: str = "") -> str:
+    ttft = f"{ttft_ms:,.0f} ms" if ttft_ms is not None else "…"
+    rate = (n_tokens / elapsed) if (elapsed > 1e-6 and n_tokens) else 0.0
+    rate_s = f"{rate:,.1f} tok/s" if n_tokens else "…"
+    if kind == "gen":
+        return ("<div class='p-status'><span class='p-live-dot'></span>"
+                f"<span class='p-status-text'>Generating · {n_tokens} tokens · TTFT {ttft} · {rate_s}</span>"
+                "<span class='p-status-hint'>press Stop to interrupt</span></div>")
+    if kind == "done":
+        return ("<div class='p-status'><span class='p-ok-dot'></span>"
+                f"<span class='p-status-text'>Done · {n_tokens} tokens · TTFT {ttft} · "
+                f"{rate:,.1f} tok/s · {elapsed:,.1f} s</span></div>")
+    if kind == "stopped":
+        return ("<div class='p-status'><span class='p-warn-dot'></span>"
+                f"<span class='p-status-text'>Stopped · {n_tokens} tokens generated · "
+                f"{elapsed:,.1f} s</span></div>")
+    if kind == "error":
+        return ("<div class='p-status'><span class='p-err-dot'></span>"
+                f"<span class='p-status-text'>Generation failed · "
+                f"{_htmllib.escape(note or 'unknown error')}</span></div>")
+    if kind == "empty":
+        return ("<div class='p-status'><span class='p-warn-dot'></span>"
+                "<span class='p-status-text'>The prompt box is empty — type something first.</span></div>")
+    return ("<div class='p-status'><span class='p-idle-dot'></span>"
+            "<span class='p-status-text'>Ready · Enter ↵ sends · Shift+Enter adds a line</span></div>")
+
+
+# -- the streaming handler ------------------------------------------------------
 def respond(message, history, max_new_tokens, temperature, top_k, top_p,
             repetition_penalty, no_repeat_ngram_size, use_anchor):
-    if not message or not message.strip():
-        yield "Please enter a prompt."
+    """Stream one continuation into the chatbot. Bound to Send, Enter, and the chips."""
+    message = (message or "").strip()
+    hist = _norm_history(history)
+
+    if not message:
+        yield (_chat_value(hist), "", _status_html("empty"), gr.update(), gr.update())
         return
 
-    prompt = (_ANCHOR + message) if use_anchor else message
-    pieces: queue.Queue = queue.Queue()
-    error: list[BaseException] = []
+    if not _ui_guard.acquire(blocking=False):
+        raise gr.Error("Pragya is already writing — press Stop or let it finish first.")
+
+    try:
+        prompt = (_ANCHOR + message) if use_anchor else message
+        stop_ev = _new_stop_event()
+        pieces: queue.Queue = queue.Queue()
+        error: list[BaseException] = []
+
+        def emit(piece: str) -> None:
+            if stop_ev.is_set():                 # raised inside generate() to abort it
+                raise _GenerationStopped()
+            pieces.put(piece)
+
+        def worker() -> None:
+            try:
+                ctx = (torch.autocast(device_type=_device.type, dtype=_dtype)
+                       if _dtype is not None else contextlib.nullcontext())
+                with _gen_lock, torch.inference_mode(), ctx:
+                    generate(
+                        _model, _tokenizer, prompt, _device,
+                        max_new_tokens=int(max_new_tokens),
+                        temperature=float(temperature),
+                        top_k=int(top_k) if top_k > 0 else None,
+                        top_p=float(top_p) if top_p < 1.0 else None,
+                        repetition_penalty=float(repetition_penalty),
+                        no_repeat_ngram_size=int(no_repeat_ngram_size),
+                        stream=False,
+                        use_cache=True,
+                        on_text=emit,
+                    )
+            except _GenerationStopped:
+                pass
+            except BaseException as exc:         # surfaced in the status line below
+                error.append(exc)
+            finally:
+                pieces.put(None)                 # end-of-stream marker
+
+        hist.append({"role": "user", "content": message})
+        hist.append({"role": "assistant", "content": "…"})
+
+        t0 = time.perf_counter()
+        ttft: float | None = None
+        text = ""
+        interrupted = False
+        threading.Thread(target=worker, daemon=True).start()
+
+        yield (_chat_value(hist), "", _status_html("gen", 0, None, 0.0),
+               gr.update(interactive=False), gr.update(visible=True))
 
-    def worker():
         try:
-            ctx = (torch.autocast(device_type=_device.type, dtype=_dtype)
-                   if _dtype is not None else contextlib.nullcontext())
-            with _gen_lock, torch.inference_mode(), ctx:
-                generate(
-                    _model, _tokenizer, prompt, _device,
-                    max_new_tokens=int(max_new_tokens),
-                    temperature=float(temperature),
-                    top_k=int(top_k) if top_k > 0 else None,
-                    top_p=float(top_p) if top_p < 1.0 else None,
-                    repetition_penalty=float(repetition_penalty),
-                    no_repeat_ngram_size=int(no_repeat_ngram_size),
-                    stream=False,
-                    use_cache=True,
-                    on_text=pieces.put,
-                )
-        except BaseException as exc:          # surfaced in the UI below
-            error.append(exc)
+            while True:
+                try:
+                    piece = pieces.get(timeout=0.2)
+                except queue.Empty:
+                    if stop_ev.is_set():
+                        interrupted = True
+                        break
+                    if ttft is None:             # still waiting for the first token
+                        yield (_chat_value(hist), "",
+                               _status_html("gen", 0, None, time.perf_counter() - t0),
+                               gr.update(), gr.update())
+                    continue
+                if piece is None:
+                    break
+                if ttft is None:
+                    ttft = (time.perf_counter() - t0) * 1000.0
+                text += piece
+                hist[-1]["content"] = text
+                yield (_chat_value(hist), "",
+                       _status_html("gen", _count(text), ttft, time.perf_counter() - t0),
+                       gr.update(), gr.update())
         finally:
-            pieces.put(None)                   # end-of-stream marker
-
-    t0 = time.perf_counter()
-    ttft = None
-    text = ""
-    threading.Thread(target=worker, daemon=True).start()
-    while True:
-        piece = pieces.get()
-        if piece is None:
-            break
-        if ttft is None:
-            ttft = (time.perf_counter() - t0) * 1000
-        text += piece
-        yield text + _stats_html(_count(text), ttft, time.perf_counter() - t0)
-
-    if error:
-        yield f"Generation failed: {error[0]}"
-        return
-    yield (text or "(no output)") + _stats_html(_count(text), ttft, time.perf_counter() - t0)
-
-
-def _count(text: str) -> int:
-    return len(_tokenizer.encode(text, add_bos=False, add_eos=False)) if text else 0
+            stop_ev.set()                        # halts the worker if this stream ends early
+
+        elapsed = time.perf_counter() - t0
+        if error:
+            note = f"{type(error[0]).__name__}: {error[0]}"
+            hist[-1]["content"] = text or "⚠ Generation failed — see the status line below."
+            final_status = _status_html("error", note=note)
+        else:
+            if not text:
+                hist[-1]["content"] = "*(stopped before any output)*" if interrupted else "*(no output)*"
+            final_status = _status_html("stopped" if interrupted else "done", _count(text), ttft, elapsed)
+        yield (_chat_value(hist), "", final_status,
+               gr.update(interactive=True), gr.update(visible=False, interactive=False))
+    finally:
+        _ui_guard.release()
+
+
+def _clear_chat():
+    """Interrupt any running generation, wait for it to wind down, then reset the UI."""
+    _request_stop()
+    with _ui_guard:                              # held by respond() while it streams
+        pass
+    return ([], "", _status_html("idle"),
+            gr.update(interactive=True), gr.update(visible=False, interactive=False))
 
 
 # -- UI -----------------------------------------------------------------------
@@ -175,83 +353,198 @@ CSS = """
 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500&display=swap');
 
 :root {
-    --p-bg: #fafaf9; --p-card: #ffffff; --p-line: #e7e5e4; --p-line-2: #d6d3d1;
-    --p-ink: #0a0a0a; --p-ink-2: #292524; --p-muted: #78716c; --p-faint: #a8a29e;
-    --p-accent: #ea580c; --p-accent-hover: #c2410c; --p-accent-soft: #fff7ed; --p-accent-line: #fed7aa;
+    --p-bg: #fafaf9; --p-card: #ffffff;
+    --p-line: #e7e5e4; --p-line-2: #d6d3d1;
+    --p-ink: #1c1917; --p-ink-2: #292524; --p-muted: #78716c; --p-faint: #a8a29e;
+    --p-accent: #ea580c; --p-accent-hover: #c2410c;
+    --p-accent-soft: #fff7ed; --p-accent-line: #fed7aa;
+    --p-shadow: 0 1px 2px rgba(28,25,23,.04), 0 16px 40px -24px rgba(28,25,23,.25);
 }
-body, .gradio-container { font-family: 'Inter', system-ui, sans-serif !important; background: var(--p-bg) !important; color: var(--p-ink) !important; }
-.gradio-container { max-width: 860px !important; margin: 0 auto !important; padding: 0 20px !important; }
-footer { display: none !important; }
 
-/* header */
-.p-nav { display: flex; align-items: center; gap: 12px; padding: 18px 0 16px; border-bottom: 1px solid var(--p-line); }
-.p-wordmark { font-weight: 700; font-size: 1.05rem; letter-spacing: -0.02em; display: flex; align-items: center; gap: 8px; }
-.p-dot { width: 8px; height: 8px; border-radius: 50%; background: var(--p-accent); }
-.p-tag { font-size: 0.72rem; font-weight: 500; color: var(--p-accent); background: var(--p-accent-soft);
-         border: 1px solid var(--p-accent-line); padding: 3px 9px; border-radius: 999px; }
-.p-nav-right { margin-left: auto; display: flex; gap: 14px; font-size: 0.78rem; color: var(--p-muted); font-family: 'IBM Plex Mono', monospace; }
-
-.p-hero { padding: 48px 0 36px; border-bottom: 1px solid var(--p-line); }
-.p-eyebrow { font-size: 0.74rem; font-weight: 600; letter-spacing: 0.08em; text-transform: uppercase; color: var(--p-accent); margin-bottom: 14px; }
-.p-hero h1 { font-size: clamp(2rem, 5vw, 2.7rem); font-weight: 700; letter-spacing: -0.03em; line-height: 1.12; margin: 0 0 14px; max-width: 580px; }
-.p-lede { font-size: 1rem; color: var(--p-muted); line-height: 1.65; max-width: 520px; margin: 0 0 26px; }
-.p-stats { display: flex; gap: 30px; flex-wrap: wrap; }
-.p-stat { display: flex; flex-direction: column; gap: 2px; }
-.p-stat b { font-size: 1.05rem; font-weight: 700; letter-spacing: -0.02em; font-variant-numeric: tabular-nums; }
-.p-stat span { font-size: 0.7rem; color: var(--p-faint); font-weight: 500; text-transform: uppercase; letter-spacing: 0.05em; }
-
-.p-section { font-size: 0.72rem; font-weight: 600; letter-spacing: 0.08em; text-transform: uppercase; color: var(--p-faint); padding: 26px 0 10px; }
-
-/* chat */
-.message-row.user-row .message, .message-row.user-row .message * { color: #ffffff !important; }
-.message-row.user-row .message { background: var(--p-accent) !important; border-radius: 12px 12px 2px 12px !important; }
-.message-row.bot-row .message, .message-row.bot-row .message * { color: #0a0a0a !important; }
-.message-row.bot-row .message code { color: #292524 !important; }
-.message-row.bot-row .message a { color: var(--p-accent) !important; }
-.message-row.bot-row .message { background: var(--p-card) !important; border: 1px solid var(--p-line) !important; border-radius: 12px 12px 12px 2px !important; box-shadow: 0 1px 3px rgba(0,0,0,0.05) !important; }
-.bot-row code, .bot-row pre { font-family: 'IBM Plex Mono', monospace !important; background: var(--p-bg) !important; border: 1px solid var(--p-line) !important; border-radius: 6px !important; font-size: 0.82rem !important; }
-.p-stats-line { font-family: 'IBM Plex Mono', monospace; font-size: 0.7rem; color: var(--p-faint); margin-top: 8px; }
-
-textarea { font-family: 'Inter', sans-serif !important; font-size: 0.95rem !important; border-radius: 8px !important; border: 1.5px solid var(--p-line-2) !important; background: var(--p-card) !important; }
-textarea:focus { border-color: var(--p-accent) !important; box-shadow: 0 0 0 3px rgba(234,88,12,0.12) !important; outline: none !important; }
-button.primary { background: var(--p-accent) !important; border: none !important; border-radius: 8px !important; font-weight: 600 !important; color: #fff !important; }
-button.primary:hover { background: var(--p-accent-hover) !important; }
-.accordion { border: 1px solid var(--p-line) !important; border-radius: 10px !important; background: var(--p-card) !important; }
-.example { font-size: 0.82rem !important; border: 1px solid var(--p-line-2) !important; border-radius: 6px !important; background: var(--p-card) !important; }
-.example:hover { border-color: var(--p-accent) !important; color: var(--p-accent) !important; background: var(--p-accent-soft) !important; }
-
-/* footer */
-.p-footer { border-top: 1px solid var(--p-line); padding: 24px 0 18px; margin-top: 14px; display: flex; justify-content: space-between; align-items: flex-start; flex-wrap: wrap; gap: 14px; }
-.p-footer-brand { font-weight: 700; font-size: 0.88rem; }
-.p-footer-copy { font-size: 0.72rem; color: var(--p-faint); margin-top: 4px; max-width: 420px; }
+* { box-sizing: border-box; }
+body, .gradio-container {
+    font-family: 'Inter', system-ui, -apple-system, 'Segoe UI', sans-serif !important;
+    background: var(--p-bg) !important;
+    color: var(--p-ink) !important;
+}
+.gradio-container {
+    max-width: 880px !important;
+    margin: 0 auto !important;
+    padding: 0 20px 48px !important;
+    background:
+        radial-gradient(900px 320px at 50% -80px, rgba(255,237,213,.9) 0%, rgba(255,237,213,0) 60%),
+        var(--p-bg) !important;
+}
+footer { display: none !important; }
+::selection { background: var(--p-accent-line); }
+
+/* ---------- header ---------- */
+.p-nav { display: flex; align-items: center; gap: 12px; padding: 20px 0 16px; border-bottom: 1px solid var(--p-line); }
+.p-wordmark { display: flex; align-items: center; gap: 9px; font-weight: 700; font-size: 1.05rem; letter-spacing: -.02em; color: var(--p-ink); }
+.p-dot { width: 9px; height: 9px; border-radius: 50%; background: var(--p-accent); box-shadow: 0 0 0 3.5px var(--p-accent-soft); }
+.p-tag { font-size: .72rem; font-weight: 600; color: var(--p-accent); background: var(--p-accent-soft); border: 1px solid var(--p-accent-line); padding: 3px 10px; border-radius: 999px; }
+.p-nav-right { margin-left: auto; display: flex; gap: 16px; font-size: .74rem; color: var(--p-faint); font-family: 'IBM Plex Mono', monospace; }
+
+.p-hero { padding: 44px 0 34px; border-bottom: 1px solid var(--p-line); margin-bottom: 6px; }
+.p-eyebrow { font-size: .72rem; font-weight: 700; letter-spacing: .1em; text-transform: uppercase; color: var(--p-accent); margin-bottom: 14px; }
+.p-hero h1 { font-size: clamp(2.05rem, 5vw, 2.8rem); font-weight: 700; letter-spacing: -.035em; line-height: 1.1; margin: 0 0 14px; color: var(--p-ink); }
+.p-grad { background: linear-gradient(92deg, #ea580c 0%, #f59e0b 100%); -webkit-background-clip: text; background-clip: text; -webkit-text-fill-color: transparent; }
+.p-lede { font-size: .98rem; color: var(--p-muted); line-height: 1.65; max-width: 540px; margin: 0 0 28px; }
+.p-stats { display: flex; flex-wrap: wrap; }
+.p-stat { display: flex; flex-direction: column; gap: 3px; padding: 0 26px; border-left: 1px solid var(--p-line); }
+.p-stat:first-child { padding-left: 0; border-left: none; }
+.p-stat b { font-size: 1.02rem; font-weight: 700; letter-spacing: -.02em; font-variant-numeric: tabular-nums; color: var(--p-ink); }
+.p-stat span { font-size: .66rem; color: var(--p-faint); font-weight: 600; text-transform: uppercase; letter-spacing: .07em; }
+
+.p-section { font-size: .7rem; font-weight: 700; letter-spacing: .1em; text-transform: uppercase; color: var(--p-faint); padding: 22px 0 10px; }
+
+/* ---------- conversation ---------- */
+.p-chat-head { align-items: center; margin-top: 10px; }
+.p-chat-title { display: flex; align-items: center; gap: 9px; font-weight: 600; font-size: .92rem; color: var(--p-ink); flex-wrap: wrap; }
+.p-chat-sub { font-weight: 400; font-size: .72rem; color: var(--p-faint); }
+.p-clear { border: 1px solid var(--p-line-2) !important; background: var(--p-card) !important; color: var(--p-muted) !important;
+           border-radius: 8px !important; font-size: .75rem !important; font-weight: 600 !important; padding: 5px 14px !important;
+           transition: all .15s ease !important; }
+.p-clear:hover { border-color: var(--p-accent) !important; color: var(--p-accent) !important; background: var(--p-accent-soft) !important; }
+
+.p-chat { border: 1px solid var(--p-line) !important; border-radius: 16px !important;
+          background: var(--p-card) !important; box-shadow: var(--p-shadow) !important; }
+
+.message-row { padding: 2px 0; }
+.message-row .message, .message.user, .message.assistant {
+    font-size: .95rem !important; line-height: 1.62 !important; padding: 10px 15px !important; margin: 2px 0 !important;
+}
+.message-row.user-row .message, .message.user {
+    background: linear-gradient(135deg, #f97316, #ea580c) !important;
+    border: 1px solid var(--p-accent-hover) !important;
+    color: #ffffff !important;
+    border-radius: 14px 14px 4px 14px !important;
+    box-shadow: 0 2px 8px -2px rgba(194,65,12,.45) !important;
+}
+.message-row.user-row .message *, .message.user * { color: #ffffff !important; }
+.message-row.bot-row .message, .message.assistant {
+    background: var(--p-card) !important;
+    border: 1px solid var(--p-line) !important;
+    color: var(--p-ink) !important;
+    border-radius: 14px 14px 14px 4px !important;
+    box-shadow: 0 1px 2px rgba(28,25,23,.05) !important;
+}
+.message-row.bot-row .message *, .message.assistant * { color: var(--p-ink) !important; }
+.message-row.bot-row .message a, .message.assistant a { color: var(--p-accent) !important; }
+.message-row.bot-row code, .message-row.bot-row pre,
+.message.assistant code, .message.assistant pre {
+    font-family: 'IBM Plex Mono', monospace !important; background: var(--p-bg) !important;
+    border: 1px solid var(--p-line) !important; border-radius: 6px !important; font-size: .82rem !important;
+}
+.message-row img { border-radius: 50% !important; object-fit: cover; }
+
+.p-chat ::-webkit-scrollbar, .p-chat *::-webkit-scrollbar { width: 7px; height: 7px; }
+.p-chat ::-webkit-scrollbar-thumb, .p-chat *::-webkit-scrollbar-thumb { background: var(--p-line-2); border-radius: 999px; }
+.p-chat ::-webkit-scrollbar-track, .p-chat *::-webkit-scrollbar-track { background: transparent; }
+
+/* ---------- status line ---------- */
+#p-status { padding: 10px 4px 0; }
+#p-status .p-status { display: flex; align-items: center; gap: 9px; flex-wrap: wrap; min-height: 20px;
+                      font-family: 'IBM Plex Mono', monospace; font-size: .72rem; color: var(--p-muted); }
+.p-live-dot, .p-idle-dot, .p-ok-dot, .p-warn-dot, .p-err-dot { flex: 0 0 auto; width: 8px; height: 8px; border-radius: 50%; }
+.p-live-dot { background: var(--p-accent); animation: p-pulse 1.1s ease-in-out infinite; }
+.p-idle-dot { background: var(--p-faint); }
+.p-ok-dot   { background: #16a34a; }
+.p-warn-dot { background: #d97706; }
+.p-err-dot  { background: #dc2626; }
+@keyframes p-pulse { 0%,100% { opacity: 1; transform: scale(1); } 50% { opacity: .3; transform: scale(.75); } }
+.p-status-hint { color: var(--p-faint); border-left: 1px solid var(--p-line); padding-left: 9px; }
+
+/* ---------- composer ---------- */
+.p-input-row { gap: 10px !important; align-items: stretch !important; }
+.p-input textarea {
+    font-family: 'Inter', sans-serif !important; font-size: .95rem !important; line-height: 1.55 !important;
+    border: 1.5px solid var(--p-line-2) !important; border-radius: 14px !important;
+    background: var(--p-card) !important; padding: 13px 16px !important;
+    box-shadow: 0 1px 2px rgba(28,25,23,.04) !important;
+    transition: border-color .15s ease, box-shadow .15s ease !important;
+}
+.p-input textarea:focus { border-color: var(--p-accent) !important; box-shadow: 0 0 0 4px rgba(234,88,12,.12) !important; }
+.p-input textarea::placeholder { color: var(--p-faint) !important; }
+
+.p-input-actions { display: flex !important; flex-direction: column !important; }
+.p-send {
+    flex: 1 1 auto !important; min-height: 64px !important;
+    background: linear-gradient(135deg, #f97316, #ea580c) !important;
+    color: #ffffff !important; border: none !important; border-radius: 14px !important;
+    font-weight: 600 !important; font-size: .95rem !important; font-family: 'Inter', sans-serif !important;
+    box-shadow: 0 2px 10px -2px rgba(234,88,12,.55) !important;
+    transition: box-shadow .15s ease, transform .08s ease !important;
+}
+.p-send:hover:not(:disabled) { box-shadow: 0 4px 16px -4px rgba(234,88,12,.7) !important; transform: translateY(-1px); }
+.p-send:active:not(:disabled) { transform: translateY(0); }
+.p-send:disabled { opacity: .55; cursor: not-allowed; box-shadow: none !important; }
+
+.p-stop {
+    width: 100% !important; margin-top: 10px !important;
+    background: #fef2f2 !important; border: 1px solid #fecaca !important; color: #b91c1c !important;
+    border-radius: 12px !important; font-size: .82rem !important; font-weight: 600 !important;
+    padding: 8px !important; transition: all .15s ease !important;
+}
+.p-stop:hover { background: #fee2e2 !important; border-color: #fca5a5 !important; }
+
+/* ---------- quick starts ---------- */
+.p-chips { gap: 8px !important; flex-wrap: wrap !important; margin-top: 2px; }
+.p-chips .form { min-width: 0 !important; margin: 0 !important; }
+.p-chip {
+    border: 1px solid var(--p-line-2) !important; background: var(--p-card) !important; color: var(--p-ink-2) !important;
+    border-radius: 999px !important; font-size: .8rem !important; font-weight: 500 !important;
+    font-family: 'Inter', sans-serif !important; padding: 6px 15px !important; transition: all .15s ease !important;
+}
+.p-chip:hover { border-color: var(--p-accent) !important; color: var(--p-accent) !important; background: var(--p-accent-soft) !important; }
+
+/* ---------- settings ---------- */
+.p-settings { border: 1px solid var(--p-line) !important; border-radius: 14px !important;
+              background: var(--p-card) !important; box-shadow: 0 1px 2px rgba(28,25,23,.04) !important; margin-top: 18px !important; }
+.p-settings .label-wrap, .p-settings summary { font-weight: 600 !important; font-size: .88rem !important; }
+input[type="range"]   { accent-color: var(--p-accent) !important; }
+input[type="checkbox"] { accent-color: var(--p-accent) !important; }
+
+/* ---------- footer ---------- */
+.p-footer { border-top: 1px solid var(--p-line); padding: 26px 0 8px; margin-top: 26px;
+            display: flex; justify-content: space-between; align-items: flex-start; flex-wrap: wrap; gap: 16px; }
+.p-footer-brand { font-weight: 700; font-size: .9rem; color: var(--p-ink); }
+.p-footer-copy { font-size: .74rem; color: var(--p-faint); margin-top: 5px; max-width: 460px; line-height: 1.6; }
 .p-badges { display: flex; gap: 6px; flex-wrap: wrap; }
-.p-badge { font-size: 0.68rem; font-family: 'IBM Plex Mono', monospace; color: var(--p-muted); background: var(--p-bg); border: 1px solid var(--p-line); padding: 3px 8px; border-radius: 4px; white-space: nowrap; }
+.p-badge { font-size: .67rem; font-family: 'IBM Plex Mono', monospace; color: var(--p-muted);
+           background: var(--p-card); border: 1px solid var(--p-line); padding: 3px 9px; border-radius: 5px; white-space: nowrap; }
 .p-badge.acc { color: var(--p-accent); background: var(--p-accent-soft); border-color: var(--p-accent-line); }
-.p-stats { font-family: 'IBM Plex Mono', monospace; font-size: 0.7rem; color: var(--p-faint); margin-top: 10px; }
+
+@media (max-width: 640px) {
+    .p-nav-right { display: none; }
+    .p-stat { padding: 0 14px !important; }
+    .p-hero { padding: 32px 0 24px; }
+}
 """
 
+_QUICK_STARTS = ("Hello", "Tell me a story.", "Create a simple story.", "Who developed you?")
+
 
 def _header_html() -> str:
     return f"""
     <div class="p-nav">
         <span class="p-wordmark"><span class="p-dot"></span>Pragya</span>
         <span class="p-tag">by Arush Kumar</span>
-        <div class="p-nav-right"><span>PyTorch</span><span>{_device_lbl}</span></div>
+        <div class="p-nav-right"><span>pytorch</span><span>{_device_lbl.lower()}</span><span>{_n_params}</span></div>
     </div>
     <div class="p-hero">
         <div class="p-eyebrow">From-scratch language model</div>
-        <h1>A model built from the ground up.</h1>
-        <p class="p-lede">A GPT-style architecture with KV-cached decoding, grouped-query attention,
-        and every weight trained from scratch. Type a prompt and watch it continue.</p>
+        <h1>A model built <span class="p-grad">from the ground up.</span></h1>
+        <p class="p-lede">A GPT-style architecture with KV-cached decoding, grouped-query attention and
+        every weight trained from scratch. Type a prompt and watch it continue — streamed token by token.</p>
         <div class="p-stats">
             <div class="p-stat"><b>{_n_params}</b><span>Parameters</span></div>
-            <div class="p-stat"><b>{_ctx_len:,}</b><span>Context</span></div>
-            <div class="p-stat"><b>{_n_layer}L &middot; {_n_head}/{_kv_heads}H</b><span>Depth &middot; heads</span></div>
+            <div class="p-stat"><b>{_ctx_len:,}</b><span>Context window</span></div>
+            <div class="p-stat"><b>{_n_layer}L &middot; {_n_head}/{_kv_heads}H</b><span>Depth &middot; Q/KV heads</span></div>
             <div class="p-stat"><b>{_pattern}</b><span>Attention</span></div>
             <div class="p-stat"><b>{_dtype_lbl}</b><span>Inference dtype</span></div>
         </div>
     </div>
-    <div class="p-section">Try it</div>
     """
 
 
@@ -259,54 +552,128 @@ def _footer_html() -> str:
     return f"""
     <div class="p-footer">
         <div>
-            <div class="p-footer-brand">Pragya &middot; <span style="font-weight:400;color:#78716c">Arush Kumar</span></div>
-            <div class="p-footer-copy">Continuation model. Earlier turns are shown in the chat but are not sent back to the model.
-            TTFT and tok/s are measured on each generation.</div>
+            <div class="p-footer-brand">Pragya &middot; <span style="font-weight:400;color:var(--p-muted)">Arush Kumar</span></div>
+            <div class="p-footer-copy">A continuation model — earlier turns stay on screen but are never re-sent
+            to the model. Enter &#8629; sends &middot; Shift+Enter adds a newline &middot; Stop interrupts mid-stream.
+            TTFT and throughput are measured on every generation.</div>
         </div>
         <div class="p-badges">
             <span class="p-badge acc">{_n_params} params</span>
             <span class="p-badge">{_device_lbl}</span>
             <span class="p-badge">{_dtype_lbl}</span>
-            <span class="p-badge">ctx {_ctx_len}</span>
+            <span class="p-badge">ctx {_ctx_len:,}</span>
             <span class="p-badge">{_pattern}</span>
+            <span class="p-badge">pytorch</span>
         </div>
     </div>
     """
 
 
-def main():
-    with gr.Blocks(title="Pragya by Arush Kumar") as demo:
+def _make_chatbot() -> gr.Chatbot:
+    """Build the Chatbot with whichever chat-value format this Gradio supports."""
+    global _MSG_FORMAT
+    supported = inspect.signature(gr.Chatbot.__init__).parameters
+    kwargs: dict[str, Any] = dict(height=460, show_label=False, elem_classes="p-chat")
+    if "type" in supported:
+        kwargs["type"] = "messages"
+    if "avatar_images" in supported and _AVATARS:
+        kwargs["avatar_images"] = _AVATARS
+    if "show_copy_button" in supported:
+        kwargs["show_copy_button"] = True
+    if "bubble_full_width" in supported:
+        kwargs["bubble_full_width"] = False
+    if "line_breaks" in supported:
+        kwargs["line_breaks"] = True
+    chatbot = gr.Chatbot(**kwargs)
+    _MSG_FORMAT = kwargs.get("type") == "messages"
+    return chatbot
+
+
+def main() -> None:
+    with gr.Blocks(
+        title="Pragya · by Arush Kumar",
+        css=CSS,   # NOTE: css/theme belong on Blocks, not launch(), on Gradio 4/5
+        theme=gr.themes.Base(
+            primary_hue=gr.themes.colors.orange,
+            neutral_hue=gr.themes.colors.stone,
+            font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
+            font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace"],
+        ),
+    ) as demo:
         gr.HTML(_header_html())
-        gr.ChatInterface(
-            fn=respond,
-            additional_inputs=[
-                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."),
-                gr.Slider(0.5, 1.0, value=1.0, step=0.01, label="Top-p", 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. 1.0 = off."),
-                gr.Slider(0, 6, value=3, step=1, label="No-repeat n-gram",
-                          info="Bans repeating any n-gram of this length. 0 = off."),
-                gr.Checkbox(value=False, label="Narrative anchor",
-                            info="Prepend a short prose passage to set the style."),
-            ],
-            additional_inputs_accordion=gr.Accordion("Generation settings", open=False),
-            chatbot=gr.Chatbot(height=440, show_label=False),
-            textbox=gr.Textbox(placeholder="Start a sentence and let Pragya continue...",
-                               show_label=False, lines=2),
-            examples=[["Hello"], ["Tell me a story."], ["Create a simple story."], ["Who developed you?"]],
-            cache_examples=False,
-        )
+
+        # -- conversation ----------------------------------------------------------
+        with gr.Row(elem_classes="p-chat-head"):
+            with gr.Column(scale=1, min_width=0):
+                gr.HTML("<div class='p-chat-title'><span class='p-dot'></span>Conversation"
+                        "<span class='p-chat-sub'>continuations only — earlier turns are not re-sent</span></div>")
+            clear_btn = gr.Button("Clear", size="sm", scale=0, min_width=0, elem_classes="p-clear")
+
+        chatbot = _make_chatbot()
+        status = gr.HTML(_status_html("idle"), elem_id="p-status")
+
+        # -- composer (textbox + Send; Stop appears only while generating) ----------
+        with gr.Row(elem_classes="p-input-row"):
+            msg = gr.Textbox(
+                placeholder="Start a sentence and let Pragya continue it…",
+                show_label=False, lines=2, max_lines=6, autofocus=True,
+                container=False, scale=1, elem_classes="p-input",
+            )
+            with gr.Column(scale=0, min_width=118, elem_classes="p-input-actions"):
+                send_btn = gr.Button("Send", variant="primary", size="lg", elem_classes="p-send")
+        stop_btn = gr.Button("Stop generating", variant="stop", visible=False, elem_classes="p-stop")
+
+        # -- quick starts ------------------------------------------------------------
+        gr.HTML("<div class='p-section'>Quick starts</div>")
+        chips: list[tuple[str, gr.Button]] = []
+        with gr.Row(elem_classes="p-chips"):
+            for example in _QUICK_STARTS:
+                chips.append((example, gr.Button(example, size="sm", scale=0, min_width=0,
+                                                 elem_classes="p-chip")))
+
+        # -- generation settings (identical knobs/defaults to before) -----------------
+        with gr.Accordion("Generation settings", open=False, elem_classes="p-settings"):
+            with gr.Row():
+                with gr.Column():
+                    max_new_tokens = gr.Slider(16, _MAX_TOKENS_CAP, value=_DEFAULT_TOKENS, step=8,
+                                               label="Max new tokens")
+                    temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
+                    top_k = gr.Slider(0, 200, value=50, step=5, label="Top-k",
+                                      info="0 disables top-k.")
+                    repetition_penalty = gr.Slider(1.0, 2.0, value=1.3, step=0.05,
+                                                   label="Repetition penalty",
+                                                   info="Count-scaled penalty. 1.0 = off.")
+                with gr.Column():
+                    top_p = gr.Slider(0.5, 1.0, value=1.0, step=0.01, label="Top-p",
+                                      info="1.0 disables nucleus sampling.")
+                    no_repeat_ngram_size = gr.Slider(0, 6, value=3, step=1, label="No-repeat n-gram",
+                                                     info="Bans repeating any n-gram of this length. 0 = off.")
+                    use_anchor = gr.Checkbox(value=False, label="Narrative anchor",
+                                             info="Prepend a short prose passage to set the style.")
+
         gr.HTML(_footer_html())
 
+        # -- wiring --------------------------------------------------------------------
+        settings = [max_new_tokens, temperature, top_k, top_p,
+                    repetition_penalty, no_repeat_ngram_size, use_anchor]
+        inputs = [msg, chatbot, *settings]
+        outputs = [chatbot, msg, status, send_btn, stop_btn]
+
+        msg.submit(respond, inputs, outputs)
+        send_btn.click(respond, inputs, outputs)
+        stop_btn.click(_request_stop, None, None, queue=False)   # immediate, bypasses the queue
+        clear_btn.click(_clear_chat, None, outputs, queue=False)
+
+        for example, chip in chips:
+            def _run_example(history, *setting_values, _example=example):
+                yield from respond(_example, history, *setting_values)
+            chip.click(_run_example, [chatbot, *settings], outputs)
+
     demo.queue().launch(
         server_name="0.0.0.0",
         server_port=int(os.environ.get("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")),
+        show_error=True,
     )