Spaces:
Sleeping
Sleeping
| import os | |
| import re | |
| import time | |
| import gradio as gr | |
| from PIL import Image | |
| from cpu_agent import ( | |
| full_agent, | |
| save_retriever_from_pdf, | |
| load_persisted_retriever, | |
| 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, | |
| ) | |
| import psutil | |
| import os | |
| def check_memory(): | |
| process = psutil.Process(os.getpid()) | |
| # RAM usage in GB | |
| return process.memory_info().rss / 1024**3 | |
| # ========================================== | |
| # 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. | |
| """ | |
| print(f"[BEFORE MODEL LOAD] RAM Usage: {check_memory():.2f} GB") | |
| model = _load_text_model("default") | |
| print(f"[AFTER MODEL LOAD] RAM Usage: {check_memory():.2f} GB") | |
| 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, | |
| ): | |
| print(f"[STREAM] RAM Usage: {check_memory():.2f} GB") | |
| 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: legal consent is required." | |
| return | |
| if pdf_file is None: | |
| 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⏳ Building vector index — please wait..." | |
| success, message = save_retriever_from_pdf(pdf_path) | |
| if success: | |
| yield f"{message}\n\n✅ Knowledge base ready." | |
| else: | |
| 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. 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, | |
| progress=gr.Progress(track_tqdm=True) | |
| ): | |
| print(f"[START] RAM Usage: {check_memory():.2f} GB") | |
| # ── Assemble clinical text ───────────────────────────────────────────── | |
| parts = [] | |
| if lab_values.strip(): | |
| parts.append(f"Lab Values:\n{lab_values.strip()}") | |
| if behaviour_changes.strip(): | |
| parts.append(f"Behaviour / Symptoms:\n{behaviour_changes.strip()}") | |
| raw_clinical_text = "\n\n".join(parts) if parts else "No clinical data provided." | |
| wsi_path = "" | |
| if wsi_image is not None: | |
| 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.", | |
| "raw_clinical_text": raw_clinical_text, | |
| "modules_queue": [], | |
| "wsi_image_path": wsi_path, | |
| "wsi_output_path": "/tmp/annotated_wsi_output.png", | |
| "module2_risk_score": "", | |
| "module3_wsi_analysis": "", | |
| "module4_progression": "", | |
| "module5_guidelines": "", | |
| "module5_raw_chunks": "", | |
| "module5_source": "", | |
| "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) | |
| print(f"[AFTER AGENT] RAM Usage: {check_memory():.2f} GB") | |
| # ── 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_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 = "" | |
| 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) | |
| mal_stat = gr.update( | |
| value=( | |
| f"Malignant: {n_mal} / {n_total} patches ({pct_str}%)\n" | |
| f"Red = Malignant · Green = Normal · Gray = Background · Yellow = Unknown" | |
| ), | |
| visible=True | |
| ) | |
| else: | |
| mal_stat = gr.update(value=wsi_summary, visible=True) | |
| except Exception: | |
| annotated_img = None | |
| 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_m5 | |
| ) | |
| # ── 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 = ( | |
| "Creatinine: 2.5 mg/dL\nCalcium: 10.5 mg/dL\nHemoglobin: 9.0 g/dL\n" | |
| "Beta-2 Microglobulin: 5.8 mg/L\nLDH: 280 U/L\n" | |
| "M-Spike: increased 1.2 g/dL over 3 months\n" | |
| "Patient: 65-year-old Male. Transplant eligible." | |
| ) | |
| EXAMPLE_BEHAVIOUR = "Dizziness, Chest Pain, Mental Confusion" | |
| # ========================================== | |
| # CUSTOM CSS — Medical Luxury Dark Theme | |
| # ========================================== | |
| 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 { | |
| --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; | |
| --font-display: 'Playfair Display', Georgia, serif; | |
| --font-data: 'IBM Plex Mono', 'Courier New', monospace; | |
| --font-body: 'IBM Plex Sans', system-ui, sans-serif; | |
| } | |
| *, *::before, *::after { box-sizing: border-box; } | |
| body, .gradio-container { | |
| background: var(--bg-void) !important; | |
| color: var(--text-ivory) !important; | |
| font-family: var(--font-body) !important; | |
| font-weight: 300 !important; | |
| } | |
| .gradio-container { | |
| max-width: 1400px !important; | |
| margin: 0 auto !important; | |
| padding: 0 24px 48px !important; | |
| } | |
| /* ── 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: 24px; | |
| position: relative; | |
| } | |
| .header-block::before { | |
| content: ''; | |
| position: absolute; | |
| bottom: -1px; left: 0; | |
| width: 80px; height: 2px; | |
| 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; | |
| font-size: 2.4rem !important; | |
| font-weight: 700 !important; | |
| color: var(--text-ivory) !important; | |
| letter-spacing: -0.02em !important; | |
| line-height: 1.15 !important; | |
| margin: 0 0 8px !important; | |
| } | |
| .gradio-container h2, .prose h2 { | |
| font-family: var(--font-display) !important; | |
| font-size: 1.35rem !important; | |
| font-weight: 600 !important; | |
| color: var(--accent-amber) !important; | |
| letter-spacing: 0.04em !important; | |
| text-transform: uppercase !important; | |
| margin: 32px 0 16px !important; | |
| padding-bottom: 8px !important; | |
| border-bottom: 1px solid var(--border-dim) !important; | |
| } | |
| .gradio-container h3, .prose h3 { | |
| font-family: var(--font-body) !important; | |
| font-size: 0.78rem !important; | |
| font-weight: 500 !important; | |
| color: var(--text-muted) !important; | |
| letter-spacing: 0.12em !important; | |
| text-transform: uppercase !important; | |
| margin: 24px 0 12px !important; | |
| } | |
| .gradio-container h4, .prose h4 { | |
| font-family: var(--font-body) !important; | |
| font-size: 0.72rem !important; | |
| font-weight: 500 !important; | |
| color: var(--accent-teal) !important; | |
| letter-spacing: 0.14em !important; | |
| text-transform: uppercase !important; | |
| margin: 20px 0 10px !important; | |
| } | |
| .gradio-container p, .prose p { | |
| color: var(--text-muted) !important; | |
| font-size: 0.88rem !important; | |
| line-height: 1.7 !important; | |
| } | |
| .gradio-container strong, .prose strong { | |
| color: var(--text-ivory) !important; | |
| font-weight: 500 !important; | |
| } | |
| /* ── Tabs ─────────────────────────────────────────────────────────────────── */ | |
| .tab-nav { | |
| background: var(--bg-deep) !important; | |
| border-bottom: 1px solid var(--border-dim) !important; | |
| padding: 0 !important; | |
| gap: 0 !important; | |
| } | |
| .tab-nav button { | |
| font-family: var(--font-body) !important; | |
| font-size: 0.75rem !important; | |
| font-weight: 500 !important; | |
| letter-spacing: 0.1em !important; | |
| text-transform: uppercase !important; | |
| color: var(--text-faint) !important; | |
| background: transparent !important; | |
| border: none !important; | |
| border-bottom: 2px solid transparent !important; | |
| padding: 16px 28px !important; | |
| transition: all 0.2s ease !important; | |
| border-radius: 0 !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; } | |
| /* ── Panels ───────────────────────────────────────────────────────────────── */ | |
| .panel, .gradio-group, .gr-group { | |
| background: var(--bg-panel) !important; | |
| border: 1px solid var(--border-dim) !important; | |
| border-radius: 4px !important; | |
| padding: 20px !important; | |
| } | |
| /* ── Labels ───────────────────────────────────────────────────────────────── */ | |
| label span, .gradio-container label { | |
| font-family: var(--font-body) !important; | |
| font-size: 0.72rem !important; | |
| font-weight: 500 !important; | |
| letter-spacing: 0.1em !important; | |
| text-transform: uppercase !important; | |
| color: var(--text-muted) !important; | |
| } | |
| /* ── Inputs ───────────────────────────────────────────────────────────────── */ | |
| textarea, input[type="text"], .gradio-textbox textarea { | |
| background: var(--bg-input) !important; | |
| border: 1px solid var(--border-dim) !important; | |
| border-radius: 3px !important; | |
| color: var(--text-ivory) !important; | |
| font-family: var(--font-data) !important; | |
| font-size: 0.82rem !important; | |
| font-weight: 300 !important; | |
| line-height: 1.6 !important; | |
| padding: 12px 14px !important; | |
| transition: border-color 0.2s ease !important; | |
| caret-color: var(--accent-amber) !important; | |
| } | |
| textarea:focus, input[type="text"]:focus { | |
| border-color: var(--border-bright) !important; | |
| outline: none !important; | |
| box-shadow: 0 0 0 1px var(--accent-amber-dim) !important; | |
| } | |
| textarea::placeholder, input::placeholder { color: var(--text-faint) !important; font-style: italic !important; } | |
| /* ── Buttons ──────────────────────────────────────────────────────────────── */ | |
| button.primary, .gr-button-primary { | |
| background: var(--accent-amber) !important; | |
| border: none !important; | |
| border-radius: 3px !important; | |
| color: var(--bg-void) !important; | |
| font-family: var(--font-body) !important; | |
| font-size: 0.78rem !important; | |
| font-weight: 600 !important; | |
| letter-spacing: 0.1em !important; | |
| text-transform: uppercase !important; | |
| padding: 14px 32px !important; | |
| transition: all 0.2s ease !important; | |
| cursor: pointer !important; | |
| } | |
| button.primary:hover, .gr-button-primary:hover { | |
| background: #d9a94d !important; | |
| transform: translateY(-1px) !important; | |
| box-shadow: 0 4px 16px rgba(200, 150, 60, 0.25) !important; | |
| } | |
| button.secondary, .gr-button-secondary { | |
| background: transparent !important; | |
| border: 1px solid var(--border-mid) !important; | |
| border-radius: 3px !important; | |
| color: var(--text-muted) !important; | |
| font-family: var(--font-body) !important; | |
| font-size: 0.75rem !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; } | |
| /* ── Checkbox ─────────────────────────────────────────────────────────────── */ | |
| .gradio-checkbox label { | |
| color: var(--text-muted) !important; | |
| font-size: 0.82rem !important; | |
| font-weight: 300 !important; | |
| letter-spacing: 0.01em !important; | |
| text-transform: none !important; | |
| line-height: 1.6 !important; | |
| } | |
| input[type="checkbox"] { accent-color: var(--accent-amber) !important; } | |
| /* ── File / Image upload ──────────────────────────────────────────────────── */ | |
| .gradio-file { | |
| background: var(--bg-input) !important; | |
| border: 1px dashed var(--border-mid) !important; | |
| border-radius: 4px !important; | |
| transition: border-color 0.2s !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 { | |
| background: var(--bg-card) !important; | |
| border: 1px solid var(--border-dim) !important; | |
| border-radius: 3px !important; | |
| padding: 10px 16px !important; | |
| } | |
| .gradio-accordion > .label-wrap span { | |
| font-family: var(--font-body) !important; | |
| font-size: 0.72rem !important; | |
| font-weight: 500 !important; | |
| letter-spacing: 0.1em !important; | |
| text-transform: uppercase !important; | |
| color: var(--text-muted) !important; | |
| } | |
| /* ── Status box ───────────────────────────────────────────────────────────── */ | |
| .status-box { | |
| background: var(--bg-card) !important; | |
| border-left: 3px solid var(--accent-teal) !important; | |
| border-top: 1px solid var(--border-dim) !important; | |
| border-right: 1px solid var(--border-dim) !important; | |
| border-bottom: 1px solid var(--border-dim) !important; | |
| border-radius: 0 3px 3px 0 !important; | |
| } | |
| hr { border: none !important; border-top: 1px solid var(--border-dim) !important; margin: 28px 0 !important; } | |
| ::-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 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 { color: var(--text-faint) !important; font-size: 0.75rem !important; } | |
| ::selection { background: var(--accent-amber-dim) !important; color: var(--text-ivory) !important; } | |
| blockquote { | |
| border-left: 3px solid var(--accent-amber-dim) !important; | |
| background: var(--bg-card) !important; | |
| margin: 0 0 12px !important; | |
| padding: 10px 16px !important; | |
| border-radius: 0 3px 3px 0 !important; | |
| } | |
| blockquote p { color: var(--text-muted) !important; font-size: 0.82rem !important; margin: 0 !important; } | |
| code { | |
| font-family: var(--font-data) !important; | |
| background: var(--bg-input) !important; | |
| color: var(--accent-teal) !important; | |
| padding: 2px 6px !important; | |
| border-radius: 2px !important; | |
| font-size: 0.82em !important; | |
| } | |
| """ | |
| # ========================================== | |
| # GRADIO UI | |
| # ========================================== | |
| 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 · Air-gapped deployment | |
| """) | |
| with gr.Tabs(): | |
| # ====================================================== | |
| # 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 becomes the evidence base for all treatment recommendations." | |
| ) | |
| gr.Markdown( | |
| "> **Supported formats:** Text-based PDF only. Scanned documents are not supported. \n" | |
| "> **Recommended sources:** ESMO, NCCN, ICMR, or your institution's local protocol." | |
| ) | |
| kb_status_box = gr.Textbox( | |
| label="Knowledge Base Status", | |
| value=get_kb_status(), | |
| interactive=False, | |
| lines=1, | |
| elem_classes=["status-box"], | |
| ) | |
| gr.Markdown("---") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| admin_pdf_upload = gr.File( | |
| label="Oncology Guideline Document", | |
| file_types=[".pdf"], | |
| file_count="single", | |
| ) | |
| 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 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") | |
| with gr.Column(scale=1): | |
| gr.Markdown(""" | |
| ### Process Overview | |
| 1. Text extracted page-by-page | |
| 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 | |
| ### Estimated Processing Time | |
| - 10-page document: 30 – 60 seconds | |
| - 50-page document: 3 – 5 minutes | |
| - 200-page document: 10 – 15 minutes | |
| *Runs once. No cloud calls.* | |
| """) | |
| 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) | |
| # ====================================================== | |
| # 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(): | |
| with gr.Column(scale=2): | |
| input_query = gr.Textbox( | |
| lines=2, | |
| label="Clinical Question", | |
| placeholder="e.g. What is the recommended treatment for this transplant-eligible patient with renal impairment?", | |
| ) | |
| with gr.Row(): | |
| input_labs = gr.Textbox( | |
| lines=9, | |
| label="Laboratory Values & Clinical Data", | |
| placeholder="Creatinine, Calcium, Haemoglobin, Beta-2 Microglobulin, LDH, M-Spike, biopsy findings, staging...", | |
| ) | |
| input_behaviour = gr.Textbox( | |
| lines=9, | |
| label="Behavioural Symptoms", | |
| placeholder="e.g. Dizziness, Chest Pain, Mental Confusion, Fatigue, Peripheral Neuropathy...", | |
| ) | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Bone Marrow Biopsy") | |
| input_wsi = gr.Image( | |
| type="pil", | |
| label="Whole Slide Image (.bmp / .png / .tiff)", | |
| height=220, | |
| ) | |
| gr.Markdown( | |
| "<small>If no image is uploaded, Module 3 generates a " | |
| "text-based WSI summary from clinical notes.</small>" | |
| ) | |
| with gr.Row(): | |
| gr.Examples( | |
| examples=[[EXAMPLE_QUERY, EXAMPLE_LABS, EXAMPLE_BEHAVIOUR, None]], | |
| 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) | |
| # ── OUTPUT ──────────────────────────────────────────────────────── | |
| gr.Markdown("---") | |
| gr.Markdown("## Analysis Output") | |
| with gr.Group(): | |
| gr.Markdown("### Orchestrator — Module Pipeline & Final Recommendation") | |
| out_modules_selected = gr.Textbox( | |
| label="Active Modules", | |
| interactive=False, | |
| lines=1, | |
| ) | |
| out_final = gr.Textbox( | |
| label="Final Recommendation (streaming · citations in [SOURCE | PAGE] format)", | |
| interactive=False, | |
| lines=10, | |
| ) | |
| gr.Markdown("---") | |
| gr.Markdown("### Specialist Module Outputs") | |
| 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) | |
| 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 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) | |
| 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) | |
| 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_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", | |
| interactive=False, lines=12, visible=False, | |
| ) | |
| submit_btn.click( | |
| fn=process_patient_data, | |
| inputs=[input_query, input_labs, input_behaviour, input_wsi], | |
| outputs=[ | |
| out_modules_selected, out_final, | |
| out_risk, grp_m2, | |
| out_wsi_text, grp_m3, | |
| out_wsi_img, out_malignancy_stat, | |
| out_prog, grp_m4, | |
| out_rag, grp_m5, | |
| out_m5_source, out_raw_chunks, | |
| ], | |
| ) | |
| demo.queue().launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, | |
| ssr_mode=False, | |
| ) |