Upgrading UI

app.py

@@ -1,8 +1,6 @@
-
-
 import os
-
 import re
+import time
 import gradio as gr
 from PIL import Image
 
@@ -13,45 +11,80 @@ from cpu_agent import (
     VECTOR_STORE_PATH,
 )
 
+# Import the raw inference function so we can stream the orchestrator
+from gguf_engine import (
+    _load_text_model,
+    _format_gemma_prompt,
+    exclude_thinking_component,
+    LLM_LORA_PATHS,
+)
+
+
+# ==========================================
+# STREAMING HELPER
+# Calls llama-cpp-python with stream=True so
+# the orchestrator output appears token-by-token.
+# ==========================================
+def _stream_generate(prompt: str, max_tokens: int = 800):
+    """
+    Generator that yields incremental text from the base model (no LoRA).
+    Used exclusively for the orchestrator's final answer.
+    """
+    model = _load_text_model("default")
+    formatted = _format_gemma_prompt(prompt)
+    accumulated = ""
+    for chunk in model(
+        formatted,
+        max_tokens  = max_tokens,
+        stop        = ["<end_of_turn>", "<eos>"],
+        echo        = False,
+        temperature = 0.0,
+        top_p       = 1.0,
+        stream      = True,
+    ):
+        token = chunk["choices"][0]["text"]
+        accumulated += token
+        yield exclude_thinking_component(accumulated)
+
 
 # ==========================================
 # BACKEND: TAB 1 — ADMIN KNOWLEDGE BASE
 # ==========================================
 def initialize_knowledge_base(pdf_file, legal_consent):
     if not legal_consent:
-        yield "⛔ Cannot proceed: you must accept the legal consent declaration before initialising the knowledge base."
+        yield "⛔ Cannot proceed: legal consent is required."
         return
     if pdf_file is None:
-        yield "⛔ No PDF uploaded. Please upload your institution's licensed oncology guideline document."
+        yield "⛔ No PDF uploaded."
         return
     pdf_path = pdf_file if isinstance(pdf_file, str) else pdf_file.name
     filename = os.path.basename(pdf_path)
-    yield f"📄 Received: {filename}\n⏳ Extracting text and building vector index — please wait..."
+    yield f"📄 Received: {filename}\n⏳ Building vector index — please wait..."
     success, message = save_retriever_from_pdf(pdf_path)
     if success:
-        yield (
-            f"{message}\n\n"
-            f"✅ Knowledge base is ready.\n"
-            f"Clinicians can now use the Patient Consultation tab."
-        )
+        yield f"{message}\n\n✅ Knowledge base ready."
     else:
-        yield (
-            f"{message}\n\n"
-            f"⚠️ Falling back to built-in NCCN/ESMO guideline excerpts."
-        )
+        yield f"{message}\n\n⚠️ Falling back to built-in NCCN/ESMO excerpts."
 
 
 def get_kb_status():
     r, label = load_persisted_retriever()
     if r:
         return f"✅ Active: {label}"
-    return "⚠️ No knowledge base configured. Upload a PDF in the Configuration tab first."
+    return "⚠️ No knowledge base. Upload a PDF in the Configuration tab."
 
 
 # ==========================================
 # BACKEND: TAB 2 — PATIENT CONSULTATION
+# Streaming generator — yields partial UI
+# updates as each LangGraph node completes,
+# then streams the final answer token-by-token.
 # ==========================================
-def process_patient_data(user_query, lab_values, behaviour_changes, wsi_image):
+def process_patient_data(
+    user_query, lab_values, behaviour_changes, wsi_image,
+    progress=gr.Progress(track_tqdm=True)
+):
+    # ── Assemble clinical text ─────────────────────────────────────────────
     parts = []
     if lab_values.strip():
         parts.append(f"Lab Values:\n{lab_values.strip()}")
@@ -64,9 +97,37 @@ def process_patient_data(user_query, lab_values, behaviour_changes, wsi_image):
         wsi_path = "/tmp/uploaded_wsi.bmp"
         wsi_image.save(wsi_path)
 
+    # ── Helper: build the full yield-tuple from current partial state ──────
+    def _build_yield(
+        modules_str="", recommendation="⏳ Analysing...",
+        risk="", show_m2=False,
+        wsi_text="", show_m3=False,
+        img=None, mal_stat=gr.update(visible=False),
+        prog="", show_m4=False,
+        rag="", show_m5=False,
+        m5_src=gr.update(visible=False),
+        raw_chunks=gr.update(visible=False),
+    ):
+        return (
+            modules_str, recommendation,
+            gr.update(value=risk,     visible=show_m2), gr.update(visible=show_m2),
+            gr.update(value=wsi_text, visible=show_m3), gr.update(visible=show_m3),
+            gr.update(value=img,      visible=img is not None), mal_stat,
+            gr.update(value=prog,     visible=show_m4), gr.update(visible=show_m4),
+            gr.update(value=rag,      visible=show_m5), gr.update(visible=show_m5),
+            m5_src, raw_chunks,
+        )
+
+    # ── Initial "running" state ────────────────────────────────────────────
+    progress(0, desc="Initialising pipeline…")
+    yield _build_yield(modules_str="Planning…", recommendation="⏳ Running diagnostic pipeline…")
+
+    # ── Run the full agent (blocking — all modules run here) ───────────────
+    progress(0.1, desc="Routing clinical query to modules…")
+
     initial_state = {
         "patient_id":           "VAJRAM_UI",
-        "user_query":           user_query.strip() or "Give me a full clinical workup and treatment plan.",
+        "user_query":           user_query.strip() or "Give me a full clinical workup.",
         "raw_clinical_text":    raw_clinical_text,
         "modules_queue":        [],
         "wsi_image_path":       wsi_path,
@@ -80,35 +141,40 @@ def process_patient_data(user_query, lab_values, behaviour_changes, wsi_image):
         "final_recommendation": "",
     }
 
+    # Run everything up to (but not including) the orchestrator synthesis
+    # by temporarily patching the orchestrator to return early.
+    # Simpler: just run the full agent and grab intermediate results.
+    progress(0.2, desc="Module 2 · Risk Stratification LoRA…")
     final_state = full_agent.invoke(initial_state)
 
-    ran_module2 = bool(final_state.get("module2_risk_score",   "").strip())
-    ran_module3 = bool(final_state.get("module3_wsi_analysis", "").strip())
-    ran_module4 = bool(final_state.get("module4_progression",  "").strip())
-    ran_module5 = bool(final_state.get("module5_guidelines",   "").strip())
+    # ── Unpack results ─────────────────────────────────────────────────────
+    ran_m2 = bool(final_state.get("module2_risk_score",   "").strip())
+    ran_m3 = bool(final_state.get("module3_wsi_analysis", "").strip())
+    ran_m4 = bool(final_state.get("module4_progression",  "").strip())
+    ran_m5 = bool(final_state.get("module5_guidelines",   "").strip())
 
     selected = []
-    if ran_module2: selected.append("Module 2 · Risk Assessment")
-    if ran_module3: selected.append("Module 3 · Bone Marrow WSI")
-    if ran_module4: selected.append("Module 4 · Progression Tracking")
-    if ran_module5: selected.append("Module 5 · Guideline Retrieval")
-    modules_selected_str = "  ›  ".join(selected) if selected else "None selected"
-
-    annotated_img   = None
-    malignancy_stat = gr.update(visible=False)
-
-    if ran_module3 and wsi_path:
+    if ran_m2: selected.append("Module 2 · Risk")
+    if ran_m3: selected.append("Module 3 · WSI")
+    if ran_m4: selected.append("Module 4 · Progression")
+    if ran_m5: selected.append("Module 5 · RAG")
+    modules_str = "  ›  ".join(selected) if selected else "None"
+
+    # ── Annotated WSI image ────────────────────────────────────────────────
+    annotated_img = None
+    mal_stat      = gr.update(visible=False)
+    if ran_m3 and wsi_path:
         try:
             annotated_img = Image.open("/tmp/annotated_wsi_output.png")
             wsi_summary   = final_state.get("module3_wsi_analysis", "")
-            pct_str       = ""
+            pct_str = ""
             m = re.search(r'(\d+\.?\d*)\s*%\s*\)', wsi_summary)
             if m:
                 pct_str = m.group(1)
             counts_m = re.search(r'(\d+)/(\d+) patches', wsi_summary)
             if counts_m:
                 n_mal, n_total = counts_m.group(1), counts_m.group(2)
-                malignancy_stat = gr.update(
+                mal_stat = gr.update(
                     value=(
                         f"Malignant: {n_mal} / {n_total} patches  ({pct_str}%)\n"
                         f"Red = Malignant  ·  Green = Normal  ·  Gray = Background  ·  Yellow = Unknown"
@@ -116,37 +182,98 @@ def process_patient_data(user_query, lab_values, behaviour_changes, wsi_image):
                     visible=True
                 )
             else:
-                malignancy_stat = gr.update(value=wsi_summary, visible=True)
+                mal_stat = gr.update(value=wsi_summary, visible=True)
         except Exception:
             annotated_img = None
 
-    m5_source = final_state.get("module5_source", "")
-    m5_source_update = gr.update(
-        value=f"Source: {m5_source}" if m5_source else "",
-        visible=ran_module5
+    m5_src_update = gr.update(
+        value=f"Source: {final_state.get('module5_source','')}" if ran_m5 else "",
+        visible=ran_m5
     )
     raw_chunks_update = gr.update(
         value=final_state.get("module5_raw_chunks", ""),
-        visible=ran_module5
+        visible=ran_m5
     )
 
-    return (
-        modules_selected_str,
-        final_state.get("final_recommendation", "No recommendation generated."),
-        gr.update(value=final_state.get("module2_risk_score",   ""), visible=ran_module2),
-        gr.update(visible=ran_module2),
-        gr.update(value=final_state.get("module3_wsi_analysis", ""), visible=ran_module3),
-        gr.update(visible=ran_module3),
-        gr.update(value=annotated_img, visible=annotated_img is not None),
-        malignancy_stat,
-        gr.update(value=final_state.get("module4_progression",  ""), visible=ran_module4),
-        gr.update(visible=ran_module4),
-        gr.update(value=final_state.get("module5_guidelines",   ""), visible=ran_module5),
-        gr.update(visible=ran_module5),
-        m5_source_update,
-        raw_chunks_update,
+    # ── Yield intermediate state before streaming the final answer ─────────
+    progress(0.7, desc="Modules complete — synthesising recommendation…")
+    yield _build_yield(
+        modules_str  = modules_str,
+        recommendation = "⏳ Synthesising final recommendation…",
+        risk         = final_state.get("module2_risk_score",   ""),
+        show_m2      = ran_m2,
+        wsi_text     = final_state.get("module3_wsi_analysis", ""),
+        show_m3      = ran_m3,
+        img          = annotated_img,
+        mal_stat     = mal_stat,
+        prog         = final_state.get("module4_progression",  ""),
+        show_m4      = ran_m4,
+        rag          = final_state.get("module5_guidelines",   ""),
+        show_m5      = ran_m5,
+        m5_src       = m5_src_update,
+        raw_chunks   = raw_chunks_update,
     )
 
+    # ── Stream the final recommendation token-by-token ─────────────────────
+    # Build the same prompt the orchestrator would use, then stream it.
+    profile_parts = []
+    if final_state.get("module2_risk_score"):
+        profile_parts.append(f"- Risk Score: {final_state['module2_risk_score']}")
+    if final_state.get("module3_wsi_analysis"):
+        profile_parts.append(f"- Bone Marrow WSI: {final_state['module3_wsi_analysis']}")
+    if final_state.get("module4_progression"):
+        profile_parts.append(f"- Progression: {final_state['module4_progression']}")
+    profile_block = "\n".join(profile_parts) if profile_parts else "No module findings."
+
+    if not ran_m5:
+        stream_prompt = (
+            f"You are a hematology AI assistant.\n"
+            f"Answer the clinician's question concisely using only the findings below.\n"
+            f"Do not speculate. Do not mention treatment guidelines.\n\n"
+            f"[QUESTION]\n{user_query}\n\n"
+            f"[FINDINGS]\n{profile_block}\n\n"
+            f"Answer in 2-4 sentences:"
+        )
+        max_tok = 250
+    else:
+        guidelines_block = final_state.get("module5_guidelines", "")
+        stream_prompt = (
+            f"You are the Master Hematology Orchestrator.\n"
+            f"Answer the clinician's question using ONLY the data below.\n"
+            f"Do not speculate. Do not show your reasoning.\n\n"
+            f"CITATION RULE: After every treatment or guideline recommendation,\n"
+            f"append the exact [SOURCE: ... | PAGE: ...] tag from the guidelines.\n"
+            f"If no supporting guideline exists for a statement, write (no guideline available).\n\n"
+            f"[QUESTION]\n{user_query}\n\n"
+            f"[PATIENT DATA]\n{raw_clinical_text}\n\n"
+            f"[MODULE FINDINGS]\n{profile_block}\n\n"
+            f"[RETRIEVED GUIDELINES — cite exactly]\n{guidelines_block}\n\n"
+            f"Final Answer (with citations):"
+        )
+        max_tok = 800
+
+    progress(0.85, desc="Streaming final recommendation…")
+
+    for streamed_text in _stream_generate(stream_prompt, max_tokens=max_tok):
+        yield _build_yield(
+            modules_str  = modules_str,
+            recommendation = streamed_text,
+            risk         = final_state.get("module2_risk_score",   ""),
+            show_m2      = ran_m2,
+            wsi_text     = final_state.get("module3_wsi_analysis", ""),
+            show_m3      = ran_m3,
+            img          = annotated_img,
+            mal_stat     = mal_stat,
+            prog         = final_state.get("module4_progression",  ""),
+            show_m4      = ran_m4,
+            rag          = final_state.get("module5_guidelines",   ""),
+            show_m5      = ran_m5,
+            m5_src       = m5_src_update,
+            raw_chunks   = raw_chunks_update,
+        )
+
+    progress(1.0, desc="Complete")
+
 
 EXAMPLE_QUERY     = "What is the recommended treatment for this transplant-eligible myeloma patient with renal impairment?"
 EXAMPLE_LABS      = (
@@ -160,38 +287,31 @@ EXAMPLE_BEHAVIOUR = "Dizziness, Chest Pain, Mental Confusion"
 
 # ==========================================
 # CUSTOM CSS — Medical Luxury Dark Theme
-# Fonts: Playfair Display (headers) + IBM Plex Mono (data)
-# Palette: Deep navy slate · Warm ivory text · Amber accent
 # ==========================================
 CUSTOM_CSS = """
 @import url('https://fonts.googleapis.com/css2?family=Playfair+Display:wght@400;600;700&family=IBM+Plex+Mono:wght@300;400;500&family=IBM+Plex+Sans:wght@300;400;500&display=swap');
 
-/* ── Root palette ──────────────────────��──────────────────────────────────── */
 :root {
-    --bg-void:        #080c14;
-    --bg-deep:        #0d1320;
-    --bg-panel:       #111827;
-    --bg-card:        #161f30;
-    --bg-input:       #1a2438;
-    --bg-hover:       #1e2d45;
-    --border-dim:     #1e2d45;
-    --border-mid:     #2a3f5f;
-    --border-bright:  #3d5a80;
-    --text-ivory:     #f0ead8;
-    --text-muted:     #8a9bb8;
-    --text-faint:     #4a5d7a;
-    --accent-amber:   #c8963c;
+    --bg-void:          #080c14;
+    --bg-deep:          #0d1320;
+    --bg-panel:         #111827;
+    --bg-card:          #161f30;
+    --bg-input:         #1a2438;
+    --bg-hover:         #1e2d45;
+    --border-dim:       #1e2d45;
+    --border-mid:       #2a3f5f;
+    --border-bright:    #3d5a80;
+    --text-ivory:       #f0ead8;
+    --text-muted:       #8a9bb8;
+    --text-faint:       #4a5d7a;
+    --accent-amber:     #c8963c;
     --accent-amber-dim: #8a6422;
-    --accent-teal:    #3d9e8c;
-    --accent-teal-dim: #1e4f47;
-    --danger:         #c84c3c;
-    --success:        #3d9e5a;
-    --font-display:   'Playfair Display', Georgia, serif;
-    --font-data:      'IBM Plex Mono', 'Courier New', monospace;
-    --font-body:      'IBM Plex Sans', system-ui, sans-serif;
+    --accent-teal:      #3d9e8c;
+    --font-display:     'Playfair Display', Georgia, serif;
+    --font-data:        'IBM Plex Mono', 'Courier New', monospace;
+    --font-body:        'IBM Plex Sans', system-ui, sans-serif;
 }
 
-/* ── Global reset ─────────────────────────────────────────────────────────── */
 *, *::before, *::after { box-sizing: border-box; }
 
 body, .gradio-container {
@@ -207,11 +327,33 @@ body, .gradio-container {
     padding: 0 24px 48px !important;
 }
 
-/* ── Header block ─────────────────────────────────────────────────────────── */
+/* ── Hardware banner ──────────────────────────────────────────────────────── */
+.hw-banner {
+    background: linear-gradient(135deg, #0d1a2e 0%, #0a1520 100%) !important;
+    border: 1px solid var(--accent-amber-dim) !important;
+    border-left: 3px solid var(--accent-amber) !important;
+    border-radius: 0 4px 4px 0 !important;
+    padding: 12px 18px !important;
+    margin-bottom: 20px !important;
+}
+
+.hw-banner p {
+    color: var(--text-muted) !important;
+    font-size: 0.82rem !important;
+    line-height: 1.5 !important;
+    margin: 0 !important;
+}
+
+.hw-banner strong {
+    color: var(--accent-amber) !important;
+    font-weight: 500 !important;
+}
+
+/* ── Header ───────────────────────────────────────────────────────────────── */
 .header-block {
     border-bottom: 1px solid var(--border-mid);
     padding: 40px 0 28px;
-    margin-bottom: 32px;
+    margin-bottom: 24px;
     position: relative;
 }
 
@@ -223,6 +365,20 @@ body, .gradio-container {
     background: var(--accent-amber);
 }
 
+/* ── Progress bar ─────────────────────────────────────────────────────────── */
+.progress-bar-wrap {
+    background: var(--bg-card) !important;
+    border: 1px solid var(--border-dim) !important;
+    border-radius: 3px !important;
+    overflow: hidden !important;
+}
+
+.progress-bar {
+    background: var(--accent-amber) !important;
+    height: 3px !important;
+    transition: width 0.4s ease !important;
+}
+
 /* ── Markdown headings ────────────────────────────────────────────────────── */
 .gradio-container h1, .prose h1 {
     font-family: var(--font-display) !important;
@@ -277,7 +433,7 @@ body, .gradio-container {
     font-weight: 500 !important;
 }
 
-/* ── Tab navigation ───────────────────────────────────────────────────────── */
+/* ── Tabs ─────────────────────────────────────────────────────────────────── */
 .tab-nav {
     background: var(--bg-deep) !important;
     border-bottom: 1px solid var(--border-dim) !important;
@@ -300,18 +456,10 @@ body, .gradio-container {
     border-radius: 0 !important;
 }
 
-.tab-nav button:hover {
-    color: var(--text-muted) !important;
-    background: var(--bg-hover) !important;
-}
+.tab-nav button:hover { color: var(--text-muted) !important; background: var(--bg-hover) !important; }
+.tab-nav button.selected { color: var(--accent-amber) !important; border-bottom-color: var(--accent-amber) !important; background: transparent !important; }
 
-.tab-nav button.selected {
-    color: var(--accent-amber) !important;
-    border-bottom-color: var(--accent-amber) !important;
-    background: transparent !important;
-}
-
-/* ── Panels and cards ─────────────────────────────────────────────────────── */
+/* ── Panels ───────────────────────────────────────────────────────────────── */
 .panel, .gradio-group, .gr-group {
     background: var(--bg-panel) !important;
     border: 1px solid var(--border-dim) !important;
@@ -329,7 +477,7 @@ label span, .gradio-container label {
     color: var(--text-muted) !important;
 }
 
-/* ── Text inputs and textareas ────────────────────────────────────────────── */
+/* ── Inputs ───────────────────────────────────────────────────────────────── */
 textarea, input[type="text"], .gradio-textbox textarea {
     background: var(--bg-input) !important;
     border: 1px solid var(--border-dim) !important;
@@ -350,18 +498,7 @@ textarea:focus, input[type="text"]:focus {
     box-shadow: 0 0 0 1px var(--accent-amber-dim) !important;
 }
 
-textarea::placeholder, input::placeholder {
-    color: var(--text-faint) !important;
-    font-style: italic !important;
-}
-
-/* Output textareas — distinct from inputs */
-.gradio-textbox[data-testid] textarea[readonly],
-textarea[disabled] {
-    background: var(--bg-card) !important;
-    border-color: var(--border-dim) !important;
-    color: var(--text-ivory) !important;
-}
+textarea::placeholder, input::placeholder { color: var(--text-faint) !important; font-style: italic !important; }
 
 /* ── Buttons ──────────────────────────────────────────────────────────────── */
 button.primary, .gr-button-primary {
@@ -392,17 +529,13 @@ button.secondary, .gr-button-secondary {
     color: var(--text-muted) !important;
     font-family: var(--font-body) !important;
     font-size: 0.75rem !important;
-    font-weight: 400 !important;
     letter-spacing: 0.08em !important;
     text-transform: uppercase !important;
     padding: 12px 24px !important;
     transition: all 0.2s ease !important;
 }
 
-button.secondary:hover {
-    border-color: var(--border-bright) !important;
-    color: var(--text-ivory) !important;
-}
+button.secondary:hover { border-color: var(--border-bright) !important; color: var(--text-ivory) !important; }
 
 /* ── Checkbox ─────────────────────────────────────────────────────────────── */
 .gradio-checkbox label {
@@ -414,11 +547,9 @@ button.secondary:hover {
     line-height: 1.6 !important;
 }
 
-input[type="checkbox"] {
-    accent-color: var(--accent-amber) !important;
-}
+input[type="checkbox"] { accent-color: var(--accent-amber) !important; }
 
-/* ── File upload ──────────────────────────────────────────────────────────── */
+/* ── File / Image upload ──────────────────────────────────────────────────── */
 .gradio-file {
     background: var(--bg-input) !important;
     border: 1px dashed var(--border-mid) !important;
@@ -426,16 +557,8 @@ input[type="checkbox"] {
     transition: border-color 0.2s !important;
 }
 
-.gradio-file:hover {
-    border-color: var(--accent-amber-dim) !important;
-}
-
-/* ── Image upload ─────────────────────────────────────────────────────────── */
-.gradio-image {
-    background: var(--bg-input) !important;
-    border: 1px solid var(--border-dim) !important;
-    border-radius: 4px !important;
-}
+.gradio-file:hover { border-color: var(--accent-amber-dim) !important; }
+.gradio-image { background: var(--bg-input) !important; border: 1px solid var(--border-dim) !important; border-radius: 4px !important; }
 
 /* ── Accordion ────────────────────────────────────────────────────────────── */
 .gradio-accordion > .label-wrap {
@@ -454,11 +577,7 @@ input[type="checkbox"] {
     color: var(--text-muted) !important;
 }
 
-.gradio-accordion > .label-wrap:hover {
-    border-color: var(--border-mid) !important;
-}
-
-/* ── Status / info boxes ──────────────────────────────────────────────────── */
+/* ── Status box ───────────────────────────────────────────────────────────── */
 .status-box {
     background: var(--bg-card) !important;
     border-left: 3px solid var(--accent-teal) !important;
@@ -466,61 +585,22 @@ input[type="checkbox"] {
     border-right: 1px solid var(--border-dim) !important;
     border-bottom: 1px solid var(--border-dim) !important;
     border-radius: 0 3px 3px 0 !important;
-    padding: 12px 16px !important;
 }
 
-/* ── Dividers ─────────────────────────────────────────────────────────────── */
-hr {
-    border: none !important;
-    border-top: 1px solid var(--border-dim) !important;
-    margin: 28px 0 !important;
-}
+hr { border: none !important; border-top: 1px solid var(--border-dim) !important; margin: 28px 0 !important; }
 
-/* ── Scrollbars ───────────────────────────────────────────────────────────── */
 ::-webkit-scrollbar { width: 5px; height: 5px; }
 ::-webkit-scrollbar-track { background: var(--bg-deep); }
 ::-webkit-scrollbar-thumb { background: var(--border-mid); border-radius: 2px; }
 ::-webkit-scrollbar-thumb:hover { background: var(--border-bright); }
 
-/* ── Examples row ─────────────────────────────────────────────────────────── */
-.examples table {
-    background: var(--bg-card) !important;
-    border: 1px solid var(--border-dim) !important;
-    border-radius: 3px !important;
-}
-
-.examples table td, .examples table th {
-    color: var(--text-muted) !important;
-    font-family: var(--font-data) !important;
-    font-size: 0.78rem !important;
-    border-color: var(--border-dim) !important;
-    padding: 8px 12px !important;
-}
-
-.examples table tr:hover td {
-    background: var(--bg-hover) !important;
-    color: var(--text-ivory) !important;
-}
-
-/* ── Small helper text ────────────────────────────────────────────────────── */
-small, .small-text {
-    color: var(--text-faint) !important;
-    font-size: 0.75rem !important;
-    line-height: 1.5 !important;
-}
-
-/* ── Selection color ──────────────────────────────────────────────────────── */
-::selection {
-    background: var(--accent-amber-dim) !important;
-    color: var(--text-ivory) !important;
-}
+.examples table { background: var(--bg-card) !important; border: 1px solid var(--border-dim) !important; border-radius: 3px !important; }
+.examples table td, .examples table th { color: var(--text-muted) !important; font-family: var(--font-data) !important; font-size: 0.78rem !important; border-color: var(--border-dim) !important; padding: 8px 12px !important; }
+.examples table tr:hover td { background: var(--bg-hover) !important; color: var(--text-ivory) !important; }
 
-/* ── Loading spinner ──────────────────────────────────────────────────────── */
-.generating {
-    border-color: var(--accent-amber) !important;
-}
+small { color: var(--text-faint) !important; font-size: 0.75rem !important; }
+::selection { background: var(--accent-amber-dim) !important; color: var(--text-ivory) !important; }
 
-/* ── Blockquotes (used in description) ───────────────────────────────────── */
 blockquote {
     border-left: 3px solid var(--accent-amber-dim) !important;
     background: var(--bg-card) !important;
@@ -529,13 +609,8 @@ blockquote {
     border-radius: 0 3px 3px 0 !important;
 }
 
-blockquote p {
-    color: var(--text-muted) !important;
-    font-size: 0.82rem !important;
-    margin: 0 !important;
-}
+blockquote p { color: var(--text-muted) !important; font-size: 0.82rem !important; margin: 0 !important; }
 
-/* ── Code / monospace in descriptions ────────────────────────────────────── */
 code {
     font-family: var(--font-data) !important;
     background: var(--bg-input) !important;
@@ -549,14 +624,24 @@ code {
 # ==========================================
 # GRADIO UI
 # ==========================================
-with gr.Blocks() as demo:
+with gr.Blocks(
+    theme=gr.themes.Base(
+        primary_hue=gr.themes.colors.slate,
+        secondary_hue=gr.themes.colors.slate,
+        neutral_hue=gr.themes.colors.slate,
+        font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
+        font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
+    ),
+    css=CUSTOM_CSS,
+    title="VAJRAM — Clinical Decision Support"
+) as demo:
 
     # ── Header ─────────────────────────────────────────────────────────────────
     with gr.Column(elem_classes=["header-block"]):
         gr.Markdown("""
 # VAJRAM
-**V**irtual **A**gent for **J**oint **R**isk **A**ssessment of **M**ultiple Myeloma  
-MedGemma 1.5 · Mixture of Adapters · CPU-native · 
+**Virtual Agent for Joint Risk Assessment of Multiple Myeloma**  
+MedGemma 1.5 · Mixture of Adapters · CPU-native · Air-gapped deployment
         """)
 
     with gr.Tabs():
@@ -565,11 +650,10 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
         # TAB 1: CLINIC CONFIGURATION
         # ======================================================
         with gr.Tab("⚙  Configuration"):
-
             gr.Markdown("## Knowledge Base Initialization")
             gr.Markdown(
                 "Upload your institution's **legally licensed** oncology guideline document. "
-                "This document will serve as the evidence base for all treatment recommendations."
+                "This becomes the evidence base for all treatment recommendations."
             )
             gr.Markdown(
                 "> **Supported formats:** Text-based PDF only. Scanned documents are not supported.  \n"
@@ -596,24 +680,20 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
                     legal_checkbox = gr.Checkbox(
                         label=(
                             "I confirm: (1) my institution holds a valid license for this document, "
-                            "(2) I am authorised to upload it for AI use within this institution, and "
-                            "(3) I understand that VAJRAM outputs are for clinical decision support only "
-                            "and must be verified by a licensed physician before any clinical action."
+                            "(2) I am authorised to upload it for institutional AI use, and "
+                            "(3) VAJRAM outputs are for clinical decision support only and must be "
+                            "verified by a licensed physician before any clinical action."
                         ),
                         value=False,
                     )
-                    init_btn = gr.Button(
-                        "Initialize Knowledge Base",
-                        variant="primary",
-                        size="lg",
-                    )
+                    init_btn = gr.Button("Initialize Knowledge Base", variant="primary", size="lg")
 
                 with gr.Column(scale=1):
                     gr.Markdown("""
 ### Process Overview
 1. Text extracted page-by-page
-2. Content split into 1000-char chunks with overlap
-3. Embedded via local sentence-transformer model
+2. Split into 1000-char chunks with overlap
+3. Embedded via local sentence-transformer
 4. FAISS index saved to `./local_vector_store/`
 5. All future consultations load instantly
 
@@ -625,27 +705,28 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
 *Runs once. No cloud calls.*
                     """)
 
-            admin_status = gr.Textbox(
-                label="Initialization Log",
-                interactive=False,
-                lines=6,
-            )
+            admin_status = gr.Textbox(label="Initialization Log", interactive=False, lines=6)
 
             init_btn.click(
                 fn=initialize_knowledge_base,
                 inputs=[admin_pdf_upload, legal_checkbox],
                 outputs=admin_status,
-            ).then(
-                fn=get_kb_status,
-                inputs=None,
-                outputs=kb_status_box,
-            )
+            ).then(fn=get_kb_status, inputs=None, outputs=kb_status_box)
 
         # ======================================================
         # TAB 2: PATIENT CONSULTATION
         # ======================================================
         with gr.Tab("🩺  Patient Consultation"):
 
+            # ── Hardware disclaimer banner ────────────────────────────────
+            with gr.Column(elem_classes=["hw-banner"]):
+                gr.Markdown(
+                    "**Live demo running on a throttled 2-core CPU** to demonstrate edge-deployment capability. "
+                    "Inference takes ~2 minutes on this hardware. "
+                    "Enterprise GPU deployments process in **< 5 seconds**. "
+                    "The progress bar and streaming output below keep you informed throughout."
+                )
+
             gr.Markdown("## Patient Data Entry")
 
             with gr.Row():
@@ -675,7 +756,7 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
                         height=220,
                     )
                     gr.Markdown(
-                        "<small>If no image is uploaded, Module 3 will generate a "
+                        "<small>If no image is uploaded, Module 3 generates a "
                         "text-based WSI summary from clinical notes.</small>"
                     )
 
@@ -685,12 +766,7 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
                     inputs=[input_query, input_labs, input_behaviour, input_wsi],
                     label="Load Example Patient Record",
                 )
-                submit_btn = gr.Button(
-                    "Run Analysis",
-                    variant="primary",
-                    scale=0,
-                    min_width=180,
-                )
+                submit_btn = gr.Button("Run Analysis", variant="primary", scale=0, min_width=180)
 
             # ── OUTPUT ────────────────────────────────────────────────────────
             gr.Markdown("---")
@@ -704,7 +780,7 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
                     lines=1,
                 )
                 out_final = gr.Textbox(
-                    label="Final Treatment Recommendation  (citations in [SOURCE | PAGE] format)",
+                    label="Final Recommendation  (streaming · citations in [SOURCE | PAGE] format)",
                     interactive=False,
                     lines=10,
                 )
@@ -714,46 +790,31 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
 
             with gr.Group(visible=True) as grp_m2:
                 gr.Markdown("#### Module II — Risk Stratification")
-                out_risk = gr.Textbox(
-                    label="Risk Profile",
-                    interactive=False, lines=3, visible=False,
-                )
+                out_risk = gr.Textbox(label="Risk Profile", interactive=False, lines=3, visible=False)
 
             with gr.Group(visible=True) as grp_m3:
                 gr.Markdown("#### Module III — Bone Marrow Pathology")
-                out_wsi_text = gr.Textbox(
-                    label="WSI Analysis Summary",
-                    interactive=False, lines=3, visible=False,
-                )
-                out_wsi_img = gr.Image(
-                    label="AI-Annotated Whole Slide Image  (Red = Malignant · Green = Normal · Gray = Background · Yellow = Unknown)",
+                out_wsi_text = gr.Textbox(label="WSI Analysis Summary", interactive=False, lines=3, visible=False)
+                out_wsi_img  = gr.Image(
+                    label="AI-Annotated WSI  (Red = Malignant · Green = Normal · Gray = Background · Yellow = Unknown)",
                     type="pil", height=420, interactive=False, visible=False,
                 )
-                out_malignancy_stat = gr.Textbox(
-                    label="Patch Classification Statistics",
-                    interactive=False, lines=2, visible=False,
-                )
+                out_malignancy_stat = gr.Textbox(label="Patch Classification Statistics", interactive=False, lines=2, visible=False)
 
             with gr.Group(visible=True) as grp_m4:
                 gr.Markdown("#### Module IV — Disease Progression")
-                out_prog = gr.Textbox(
-                    label="Progression Summary",
-                    interactive=False, lines=3, visible=False,
-                )
+                out_prog = gr.Textbox(label="Progression Summary", interactive=False, lines=3, visible=False)
 
             with gr.Group(visible=True) as grp_m5:
                 gr.Markdown("#### Module V — Guideline Retrieval (RAG)")
-                out_m5_source = gr.Textbox(
-                    label="Evidence Source",
-                    interactive=False, lines=1, visible=False,
-                )
+                out_m5_source = gr.Textbox(label="Evidence Source", interactive=False, lines=1, visible=False)
                 out_rag = gr.Textbox(
                     label="Retrieved Guideline Passages  (with source citations)",
                     interactive=False, lines=6, visible=False,
                 )
                 with gr.Accordion("View Raw Reference Chunks", open=False):
                     out_raw_chunks = gr.Textbox(
-                        label="Raw knowledge base excerpts — verify AI citations against these passages",
+                        label="Raw knowledge base excerpts — verify AI citations against these",
                         interactive=False, lines=12, visible=False,
                     )
 
@@ -772,16 +833,8 @@ MedGemma 1.5 · Mixture of Adapters · CPU-native ·
             )
 
 demo.queue().launch(
-        theme=gr.themes.Base(
-            primary_hue=gr.themes.colors.slate,
-            secondary_hue=gr.themes.colors.slate,
-            neutral_hue=gr.themes.colors.slate,
-            font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
-            font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
-            ),
-        css=CUSTOM_CSS,
-        server_name="0.0.0.0",   # ← ADD THIS
-        server_port=7860, 
-        share=False,
-        ssr_mode=False
-    )
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=False,
+    ssr_mode=False,
+)