Spaces:
Runtime error
Runtime error
| """ | |
| ChatSpatial Engine — Gemma 4 Spatial Omics | |
| Multimodal spatial transcriptomics analysis powered by Gemma 4 31B + LoRA. | |
| """ | |
| import torch | |
| import gradio as gr | |
| from PIL import Image | |
| import os | |
| from visualizations import generate_all_plots, fig_to_pil | |
| # --------------------------------------------------------------------------- | |
| # Model Loading (runs once at startup) | |
| # --------------------------------------------------------------------------- | |
| print("Loading model... This may take a few minutes on first boot.") | |
| from unsloth import FastVisionModel | |
| model, tokenizer = FastVisionModel.from_pretrained( | |
| "arka2696/gemma-4-spatial-omics-lora-v3", | |
| load_in_4bit=True, | |
| ) | |
| FastVisionModel.for_inference(model) | |
| # Set chat template for multimodal messages | |
| CHAT_TEMPLATE = ( | |
| "{% for message in messages %}" | |
| "{% if message['role'] == 'user' %}{{ '<start_of_turn>user\\n' }}" | |
| "{% elif message['role'] == 'model' %}{{ '<start_of_turn>model\\n' }}" | |
| "{% endif %}" | |
| "{% if message['content'] is string %}{{ message['content'] }}" | |
| "{% else %}{% for block in message['content'] %}" | |
| "{% if block['type'] == 'image' %}{{ '<|image|>' }}" | |
| "{% elif block['type'] == 'text' %}{{ block['text'] }}" | |
| "{% endif %}{% endfor %}{% endif %}" | |
| "{{ '<end_of_turn>\\n' }}" | |
| "{% endfor %}" | |
| ) | |
| tokenizer.chat_template = CHAT_TEMPLATE | |
| SYSTEM_PROMPT = ( | |
| "You are ChatSpatial, an expert spatial transcriptomics analyst. " | |
| "You analyze H&E-stained tissue microscopy images and predict gene expression " | |
| "patterns based on cellular morphology, tissue architecture, and spatial context. " | |
| "Always reason step-by-step about what you observe in the tissue before making predictions. " | |
| "Use <|think|> tags for your internal reasoning." | |
| ) | |
| print("Model loaded successfully!") | |
| # --------------------------------------------------------------------------- | |
| # Inference | |
| # --------------------------------------------------------------------------- | |
| def _load_phoenix(): | |
| """Load Phoenix engine (optional — fails gracefully).""" | |
| try: | |
| from phoenix_engine import PhoenixEngine | |
| from model_loader import download_phoenix_weights | |
| model_dir = "phoenix_weights" | |
| if not os.path.exists(model_dir): | |
| download_phoenix_weights(model_dir) | |
| engine = PhoenixEngine(model_dir=model_dir, device="cpu", num_samples=3) | |
| engine._load() | |
| print("Phoenix engine loaded!", flush=True) | |
| return engine | |
| except Exception as e: | |
| print(f"Phoenix not available: {e}", flush=True) | |
| return None | |
| phoenix_engine = _load_phoenix() | |
| def analyze(image: Image.Image, question: str, h5ad_file=None): | |
| """Run inference on an image + question. Returns (analysis_md, bar_img, radar_img, heatmap_img).""" | |
| if image is None: | |
| return "Please upload an H&E tissue image to analyze.", None, None, None | |
| if not question.strip(): | |
| question = "Analyze this tissue image. What cell types and gene expression patterns do you observe?" | |
| prompt = question | |
| bar_img, radar_img, heatmap_img = None, None, None | |
| # Phoenix prediction (if available + image provided) | |
| if phoenix_engine is not None: | |
| try: | |
| phoenix_result = phoenix_engine.predict(image) | |
| if phoenix_result: | |
| phoenix_text = phoenix_result.get("summary_text", "") | |
| prompt += ( | |
| "\n\nPhoenix quantitative expression prediction results:\n" | |
| + phoenix_text | |
| + "\n\nIntegrate these quantitative results with your morphological " | |
| "observations. Add biological interpretation beyond restating numbers." | |
| ) | |
| plots = generate_all_plots(phoenix_result) | |
| if "bar_chart" in plots: | |
| bar_img = fig_to_pil(plots["bar_chart"]) | |
| if "radar" in plots: | |
| radar_img = fig_to_pil(plots["radar"]) | |
| if "heatmap" in plots: | |
| heatmap_img = fig_to_pil(plots["heatmap"]) | |
| except Exception as e: | |
| print(f"Phoenix error: {e}", flush=True) | |
| # Build messages with image object directly (Unsloth's expected format) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": image}, | |
| {"type": "text", "text": f"{SYSTEM_PROMPT}\n\n{prompt}"}, | |
| ], | |
| } | |
| ] | |
| # Tokenize (Unsloth handles image processing internally) | |
| inputs = tokenizer.apply_chat_template( | |
| messages, | |
| tokenize=True, | |
| add_generation_prompt=True, | |
| return_tensors="pt", | |
| return_dict=True, | |
| ) | |
| inputs = {k: v.to(model.device) for k, v in inputs.items()} | |
| # Generate | |
| with torch.inference_mode(): | |
| output_ids = model.generate( | |
| **inputs, | |
| max_new_tokens=1024, | |
| temperature=0.6, | |
| top_p=0.9, | |
| do_sample=True, | |
| use_cache=True, | |
| ) | |
| # Decode only new tokens | |
| generated = output_ids[0][inputs["input_ids"].shape[1]:] | |
| response = tokenizer.decode(generated, skip_special_tokens=True) | |
| # Format reasoning traces nicely | |
| response = format_reasoning(response) | |
| return response, bar_img, radar_img, heatmap_img | |
| def format_reasoning(text: str) -> str: | |
| """Format <|think|> blocks as collapsible markdown sections.""" | |
| import re | |
| # Extract think blocks and format as details/summary | |
| def replace_think(match): | |
| reasoning = match.group(1).strip() | |
| return ( | |
| f"\n<details>\n<summary>💭 <b>Reasoning Trace</b></summary>\n\n" | |
| f"{reasoning}\n\n</details>\n\n" | |
| ) | |
| text = re.sub( | |
| r"<\|think\|>(.*?)<\|/think\|>", | |
| replace_think, | |
| text, | |
| flags=re.DOTALL, | |
| ) | |
| return text.strip() | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| /* ===== Design Tokens ===== */ | |
| :root { | |
| --bg-primary: #0f172a; | |
| --bg-secondary: #1e293b; | |
| --bg-card: #1e293b; | |
| --bg-card-hover: #253349; | |
| --border-subtle: rgba(255, 255, 255, 0.06); | |
| --border-active: rgba(59, 130, 246, 0.4); | |
| --accent-blue: #3b82f6; | |
| --accent-blue-dim: rgba(59, 130, 246, 0.15); | |
| --accent-green: #10b981; | |
| --accent-green-dim: rgba(16, 185, 129, 0.15); | |
| --accent-amber: #f59e0b; | |
| --accent-amber-dim: rgba(245, 158, 11, 0.15); | |
| --text-primary: #f1f5f9; | |
| --text-secondary: #94a3b8; | |
| --text-muted: #64748b; | |
| --font-sans: -apple-system, BlinkMacSystemFont, 'Segoe UI', Inter, Roboto, sans-serif; | |
| --font-mono: 'JetBrains Mono', 'Fira Code', 'SF Mono', Consolas, monospace; | |
| --radius-sm: 6px; | |
| --radius-md: 10px; | |
| --radius-lg: 14px; | |
| --shadow-card: 0 1px 3px rgba(0,0,0,0.3), 0 1px 2px rgba(0,0,0,0.2); | |
| --shadow-elevated: 0 4px 12px rgba(0,0,0,0.4); | |
| } | |
| /* ===== Global ===== */ | |
| .gradio-container { | |
| background: var(--bg-primary) !important; | |
| font-family: var(--font-sans) !important; | |
| max-width: 1440px !important; | |
| color: var(--text-primary) !important; | |
| } | |
| .gradio-container .gr-block, | |
| .gradio-container .gr-box, | |
| .gradio-container .gr-panel { | |
| background: transparent !important; | |
| } | |
| footer { display: none !important; } | |
| /* ===== Top Navbar ===== */ | |
| .navbar { | |
| display: flex; | |
| align-items: center; | |
| justify-content: space-between; | |
| padding: 12px 24px; | |
| background: var(--bg-secondary); | |
| border-bottom: 1px solid var(--border-subtle); | |
| border-radius: var(--radius-lg); | |
| margin-bottom: 16px; | |
| } | |
| .navbar-brand { | |
| display: flex; | |
| align-items: center; | |
| gap: 10px; | |
| font-size: 1.1rem; | |
| font-weight: 700; | |
| color: var(--text-primary); | |
| letter-spacing: -0.02em; | |
| } | |
| .navbar-brand .dot { | |
| color: var(--accent-green); | |
| font-size: 1.3rem; | |
| line-height: 1; | |
| } | |
| .navbar-right { | |
| display: flex; | |
| align-items: center; | |
| gap: 8px; | |
| } | |
| .nav-badge { | |
| display: inline-flex; | |
| align-items: center; | |
| gap: 5px; | |
| padding: 4px 10px; | |
| border-radius: 20px; | |
| font-size: 0.72rem; | |
| font-weight: 600; | |
| font-family: var(--font-mono); | |
| letter-spacing: 0.02em; | |
| text-transform: uppercase; | |
| } | |
| .nav-badge.blue { | |
| background: var(--accent-blue-dim); | |
| color: var(--accent-blue); | |
| border: 1px solid rgba(59, 130, 246, 0.25); | |
| } | |
| .nav-badge.green { | |
| background: var(--accent-green-dim); | |
| color: var(--accent-green); | |
| border: 1px solid rgba(16, 185, 129, 0.25); | |
| } | |
| .nav-badge.amber { | |
| background: var(--accent-amber-dim); | |
| color: var(--accent-amber); | |
| border: 1px solid rgba(245, 158, 11, 0.25); | |
| } | |
| .status-indicator { | |
| display: inline-flex; | |
| align-items: center; | |
| gap: 6px; | |
| padding: 4px 12px; | |
| border-radius: 20px; | |
| font-size: 0.75rem; | |
| font-weight: 500; | |
| background: var(--accent-green-dim); | |
| color: var(--accent-green); | |
| border: 1px solid rgba(16, 185, 129, 0.3); | |
| } | |
| .status-indicator .status-dot { | |
| width: 7px; | |
| height: 7px; | |
| background: var(--accent-green); | |
| border-radius: 50%; | |
| animation: pulse 2s ease-in-out infinite; | |
| } | |
| @keyframes pulse { | |
| 0%, 100% { opacity: 1; } | |
| 50% { opacity: 0.5; } | |
| } | |
| /* ===== Card Containers ===== */ | |
| .ui-card { | |
| background: var(--bg-card) !important; | |
| border: 1px solid var(--border-subtle) !important; | |
| border-radius: var(--radius-lg) !important; | |
| padding: 20px !important; | |
| box-shadow: var(--shadow-card) !important; | |
| transition: border-color 0.2s ease !important; | |
| } | |
| .ui-card:hover { | |
| border-color: rgba(255, 255, 255, 0.1) !important; | |
| } | |
| .card-header { | |
| display: flex; | |
| align-items: center; | |
| gap: 8px; | |
| margin-bottom: 12px; | |
| font-size: 0.82rem; | |
| font-weight: 600; | |
| color: var(--text-secondary); | |
| text-transform: uppercase; | |
| letter-spacing: 0.05em; | |
| } | |
| .card-header .card-icon { | |
| font-size: 0.9rem; | |
| } | |
| /* ===== Inputs Styling ===== */ | |
| .gradio-container textarea, | |
| .gradio-container input[type="text"] { | |
| background: var(--bg-primary) !important; | |
| border: 1px solid var(--border-subtle) !important; | |
| border-radius: var(--radius-md) !important; | |
| color: var(--text-primary) !important; | |
| font-family: var(--font-sans) !important; | |
| font-size: 0.9rem !important; | |
| transition: border-color 0.2s ease !important; | |
| } | |
| .gradio-container textarea:focus, | |
| .gradio-container input[type="text"]:focus { | |
| border-color: var(--accent-blue) !important; | |
| box-shadow: 0 0 0 3px var(--accent-blue-dim) !important; | |
| } | |
| /* ===== Primary Button ===== */ | |
| #run-btn { | |
| background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%) !important; | |
| border: none !important; | |
| border-radius: var(--radius-md) !important; | |
| color: #ffffff !important; | |
| font-weight: 600 !important; | |
| font-size: 0.92rem !important; | |
| padding: 12px 28px !important; | |
| letter-spacing: 0.01em !important; | |
| transition: all 0.15s ease !important; | |
| box-shadow: 0 2px 8px rgba(59, 130, 246, 0.25) !important; | |
| } | |
| #run-btn:hover { | |
| transform: translateY(-1px) !important; | |
| box-shadow: 0 6px 20px rgba(59, 130, 246, 0.35) !important; | |
| } | |
| #run-btn:active { | |
| transform: translateY(0) !important; | |
| } | |
| /* ===== Tabs ===== */ | |
| .gradio-container .tabs { | |
| background: transparent !important; | |
| } | |
| .gradio-container .tab-nav { | |
| border-bottom: 1px solid var(--border-subtle) !important; | |
| background: transparent !important; | |
| gap: 0 !important; | |
| } | |
| .gradio-container .tab-nav button { | |
| background: transparent !important; | |
| border: none !important; | |
| border-bottom: 2px solid transparent !important; | |
| color: var(--text-muted) !important; | |
| font-weight: 500 !important; | |
| font-size: 0.85rem !important; | |
| padding: 10px 18px !important; | |
| transition: all 0.15s ease !important; | |
| } | |
| .gradio-container .tab-nav button:hover { | |
| color: var(--text-secondary) !important; | |
| } | |
| .gradio-container .tab-nav button.selected { | |
| color: var(--accent-blue) !important; | |
| border-bottom-color: var(--accent-blue) !important; | |
| } | |
| /* ===== Output Markdown ===== */ | |
| .output-markdown { | |
| background: var(--bg-primary) !important; | |
| border: 1px solid var(--border-subtle) !important; | |
| border-radius: var(--radius-md) !important; | |
| padding: 24px !important; | |
| color: var(--text-primary) !important; | |
| font-size: 0.9rem !important; | |
| line-height: 1.7 !important; | |
| min-height: 300px !important; | |
| } | |
| .output-markdown code { | |
| font-family: var(--font-mono) !important; | |
| background: rgba(59, 130, 246, 0.1) !important; | |
| padding: 2px 6px !important; | |
| border-radius: 4px !important; | |
| font-size: 0.82rem !important; | |
| } | |
| .output-markdown details { | |
| background: rgba(16, 185, 129, 0.05) !important; | |
| border: 1px solid rgba(16, 185, 129, 0.15) !important; | |
| border-radius: var(--radius-sm) !important; | |
| padding: 12px !important; | |
| margin: 12px 0 !important; | |
| } | |
| .output-markdown details summary { | |
| cursor: pointer; | |
| color: var(--accent-green) !important; | |
| font-weight: 600 !important; | |
| } | |
| /* ===== Image Upload ===== */ | |
| .gradio-container .image-container, | |
| .gradio-container .upload-container { | |
| border: 1px dashed var(--border-subtle) !important; | |
| border-radius: var(--radius-md) !important; | |
| background: var(--bg-primary) !important; | |
| transition: border-color 0.2s ease !important; | |
| } | |
| .gradio-container .image-container:hover, | |
| .gradio-container .upload-container:hover { | |
| border-color: var(--accent-blue) !important; | |
| } | |
| /* ===== Plot Outputs ===== */ | |
| .gradio-container .plot-container, | |
| .gradio-container .image-preview { | |
| border-radius: var(--radius-md) !important; | |
| overflow: hidden !important; | |
| } | |
| /* ===== Section Label ===== */ | |
| .section-label { | |
| font-size: 0.75rem; | |
| font-weight: 600; | |
| color: var(--text-muted); | |
| text-transform: uppercase; | |
| letter-spacing: 0.06em; | |
| margin-bottom: 10px; | |
| padding-left: 2px; | |
| } | |
| /* ===== Footer ===== */ | |
| .app-footer { | |
| text-align: center; | |
| padding: 20px 16px; | |
| margin-top: 24px; | |
| border-top: 1px solid var(--border-subtle); | |
| } | |
| .app-footer .footer-text { | |
| color: var(--text-muted); | |
| font-size: 0.78rem; | |
| letter-spacing: 0.01em; | |
| } | |
| .footer-badges { | |
| display: flex; | |
| justify-content: center; | |
| gap: 6px; | |
| margin-top: 8px; | |
| flex-wrap: wrap; | |
| } | |
| .footer-badge { | |
| display: inline-flex; | |
| align-items: center; | |
| padding: 3px 8px; | |
| border-radius: 4px; | |
| font-size: 0.68rem; | |
| font-family: var(--font-mono); | |
| font-weight: 500; | |
| background: rgba(255, 255, 255, 0.04); | |
| color: var(--text-muted); | |
| border: 1px solid var(--border-subtle); | |
| } | |
| /* ===== Dropdown ===== */ | |
| .gradio-container .dropdown-container, | |
| .gradio-container select { | |
| background: var(--bg-primary) !important; | |
| border: 1px solid var(--border-subtle) !important; | |
| border-radius: var(--radius-md) !important; | |
| color: var(--text-primary) !important; | |
| } | |
| /* ===== Label styling ===== */ | |
| .gradio-container label { | |
| color: var(--text-secondary) !important; | |
| font-size: 0.82rem !important; | |
| font-weight: 500 !important; | |
| } | |
| /* ===== Responsive ===== */ | |
| @media (max-width: 768px) { | |
| .navbar { | |
| flex-direction: column; | |
| gap: 10px; | |
| text-align: center; | |
| } | |
| .navbar-right { | |
| flex-wrap: wrap; | |
| justify-content: center; | |
| } | |
| } | |
| """ | |
| EXAMPLES = [ | |
| "Analyze this tissue patch. What cell types do you observe based on morphology?", | |
| "Predict the expression levels of CD8A, EPCAM, and COL1A1 in this region.", | |
| "Is this region likely tumor, stroma, or immune-infiltrated? Explain your reasoning.", | |
| "What spatial gene expression patterns would you expect in this tissue architecture?", | |
| "Identify the dominant cell population and predict marker gene expression.", | |
| ] | |
| # Determine component status for navbar | |
| _model_status = "Ready" if model is not None else "Loading" | |
| _phoenix_status = "Active" if phoenix_engine is not None else "Inactive" | |
| with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="ChatSpatial Engine") as demo: | |
| # ===== Top Navbar ===== | |
| gr.HTML(f""" | |
| <div class="navbar"> | |
| <div class="navbar-brand"> | |
| <span class="dot">◉</span> | |
| ChatSpatial Engine | |
| </div> | |
| <div class="navbar-right"> | |
| <span class="nav-badge blue">Gemma 4 31B</span> | |
| <span class="nav-badge green">Phoenix</span> | |
| <span class="nav-badge amber">QLoRA</span> | |
| <div class="status-indicator"> | |
| <span class="status-dot"></span> | |
| {_model_status} | |
| </div> | |
| </div> | |
| </div> | |
| """) | |
| # ===== Main Layout ===== | |
| with gr.Row(equal_height=False): | |
| # ----- Left Sidebar (Inputs) ----- | |
| with gr.Column(scale=4): | |
| # Image Upload Card | |
| gr.HTML(""" | |
| <div class="card-header"> | |
| <span class="card-icon">🔬</span> H&E Tissue Patch | |
| </div> | |
| """) | |
| image_input = gr.Image( | |
| type="pil", | |
| label="Drop or click to upload tissue image", | |
| height=260, | |
| elem_classes=["ui-card"], | |
| ) | |
| gr.HTML("<div style='height: 12px;'></div>") | |
| # Question Card | |
| gr.HTML(""" | |
| <div class="card-header"> | |
| <span class="card-icon">❓</span> Analysis Query | |
| </div> | |
| """) | |
| question_input = gr.Textbox( | |
| label="Question", | |
| placeholder="Ask about cell types, gene expression, tissue architecture...", | |
| lines=3, | |
| ) | |
| example_dropdown = gr.Dropdown( | |
| choices=EXAMPLES, | |
| label="Example prompts", | |
| interactive=True, | |
| ) | |
| gr.HTML("<div style='height: 12px;'></div>") | |
| # h5ad optional upload | |
| gr.HTML(""" | |
| <div class="card-header"> | |
| <span class="card-icon">📊</span> Spatial Data (optional) | |
| </div> | |
| """) | |
| h5ad_input = gr.File( | |
| label=".h5ad spatial matrix", | |
| file_types=[".h5ad"], | |
| ) | |
| gr.HTML("<div style='height: 16px;'></div>") | |
| # Run Button | |
| analyze_btn = gr.Button( | |
| "▶ Run Analysis", | |
| variant="primary", | |
| elem_id="run-btn", | |
| ) | |
| # ----- Right Panel (Outputs) ----- | |
| with gr.Column(scale=6): | |
| with gr.Tabs(): | |
| # Tab: Analysis | |
| with gr.Tab("Analysis"): | |
| output_md = gr.Markdown( | |
| value="*Upload an H&E tissue image and click* **Run Analysis** *to begin.*", | |
| elem_classes=["output-markdown"], | |
| ) | |
| # Tab: Expression Profile | |
| with gr.Tab("Expression Profile"): | |
| gr.HTML('<div class="section-label">Gene Expression Visualization</div>') | |
| with gr.Row(): | |
| bar_output = gr.Image(label="Top Expressed Genes", height=320) | |
| radar_output = gr.Image(label="Tissue Composition Radar", height=320) | |
| heatmap_output = gr.Image(label="Marker Gene Heatmap", height=200) | |
| # Tab: About | |
| with gr.Tab("About"): | |
| gr.Markdown(""" | |
| ### Methodology | |
| **ChatSpatial Engine** is a multimodal spatial transcriptomics analysis system built on: | |
| - **Gemma 4 31B Dense** fine-tuned with QLoRA (4-bit quantization) on the STimage-1K4M dataset | |
| - **Phoenix Flow-Matching** for quantitative gene expression prediction from morphology | |
| - **Chain-of-thought reasoning** via `<|think|>` traces for interpretable biological analysis | |
| **Training Data**: 4M+ spatial transcriptomics spots across 1,171 H&E-stained tissue slides | |
| (Visium, ST, VisiumHD technologies) with matched gene expression profiles. | |
| **Marker Panel**: CD8A, CD8B, CD3D, CD4, MS4A1, CD19, CD68, CD163, PTPRC, EPCAM, KRT18, | |
| MKI67, COL1A1, VIM, ACTA2, VEGFA, PDCD1 | |
| **Pipeline**: Image patch (224x224) -> Vision encoder (1120 tokens max) -> Gemma 4 reasoning | |
| -> Expression prediction + biological interpretation | |
| --- | |
| *Built for the Gemma 4 Good Hackathon (Kaggle)* | |
| """) | |
| # ===== Wire Events ===== | |
| analyze_btn.click( | |
| fn=analyze, | |
| inputs=[image_input, question_input, h5ad_input], | |
| outputs=[output_md, bar_output, radar_output, heatmap_output], | |
| ) | |
| example_dropdown.change( | |
| fn=lambda x: x, | |
| inputs=example_dropdown, | |
| outputs=question_input, | |
| ) | |
| # ===== Footer ===== | |
| gr.HTML(""" | |
| <div class="app-footer"> | |
| <div class="footer-text">Powered by</div> | |
| <div class="footer-badges"> | |
| <span class="footer-badge">Gemma 4 31B</span> | |
| <span class="footer-badge">STimage-1K4M</span> | |
| <span class="footer-badge">Unsloth QLoRA</span> | |
| <span class="footer-badge">Phoenix Flow-Matching</span> | |
| <span class="footer-badge">HuggingFace Spaces</span> | |
| </div> | |
| </div> | |
| """) | |
| # --------------------------------------------------------------------------- | |
| # Launch | |
| # --------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |