Update app with casual examples, CPU mode, state persistence

app.py

@@ -9,13 +9,24 @@ MODEL_ID = "Janushi/ClinicalDistill-Gemma-1B"
 tokenizer = None
 model = None
 
-INSTRUCTION = """Extract symptoms from the clinical note and return JSON with this exact format:
+INSTRUCTION = """You are a clinical NLP model. Extract ONLY medical symptoms from the clinical note.
+
+Return JSON in this exact format:
 {
   "symptoms": ["symptom1", "symptom2"],
   "duration": ["duration1", "duration2"],
   "severity": ["severity1", "severity2"],
   "urgent": true/false
-}"""
+}
+
+Rules:
+- symptoms: ONLY medical symptoms (fever, back pain, headache, nausea, cough, dizziness). NOT observations, context, or descriptions like "seems okay", "a little cranky", "not sure"
+- duration: how long each symptom has lasted. Use "unspecified" if not mentioned
+- severity: how severe each symptom is. Use "unspecified" if not clearly stated — do NOT guess severity
+- urgent=true ONLY for: chest pain, difficulty breathing, stroke symptoms (slurred speech, facial drooping, arm weakness), severe bleeding, loss of consciousness
+- urgent=false for: back pain, headache, nausea, fever, diarrhea, fatigue, sneezing, runny nose, cough, irritability, dizziness, stomach ache
+- Never duplicate symptoms
+- All arrays must be the same length"""
 
 EXAMPLES = [
     ["been feeling off for a few days, chest feels weird and i get tired just walking to the kitchen"],
@@ -59,6 +70,28 @@ CSS = """
 footer { display: none !important; }
 """
 
+# Phrases that are NOT valid medical symptoms
+NON_SYMPTOM_PHRASES = [
+    "seems okay", "seems fine", "otherwise fine", "no fever", "a little",
+    "otherwise", "seems", "appears", "looks", "none", "normal", "okay",
+    "fine", "not sure", "cranky", "irritable", "fussy", "acting up",
+    "feeling off", "feeling drained", "feeling tired", "feeling weak"
+]
+
+def is_valid_symptom(s: str) -> bool:
+    s_lower = s.lower().strip()
+    # Too long to be a real symptom (more than 5 words)
+    if len(s_lower.split()) > 5:
+        return False
+    # Contains non-symptom phrases
+    if any(phrase in s_lower for phrase in NON_SYMPTOM_PHRASES):
+        return False
+    # Too short to be meaningful
+    if len(s_lower) < 3:
+        return False
+    return True
+
+
 def load_model():
     global tokenizer, model
     if model is not None:
@@ -67,10 +100,11 @@ def load_model():
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         torch_dtype=torch.float32,  # float32 for CPU
-        device_map="cpu"
+        device_map="cpu",
     )
     model.eval()
 
+
 def build_prompt(clinical_note: str) -> str:
     return (
         f"<instruction>\n{INSTRUCTION}\n</instruction>\n\n"
@@ -78,36 +112,82 @@ def build_prompt(clinical_note: str) -> str:
         f"<output>\n"
     )
 
+
 def parse_output(raw: str) -> dict:
     raw = raw.split("</output>")[0].strip()
     match = re.search(r"\{.*\}", raw, re.DOTALL)
     if match:
         raw = match.group(0)
     result = json.loads(raw)
+
     for key in ("symptoms", "duration", "severity"):
         if key not in result or not isinstance(result[key], list):
             result[key] = []
     if "urgent" not in result:
         result["urgent"] = False
-    n = max(len(result["symptoms"]), len(result["duration"]), len(result["severity"]), 1)
+
+    # Deduplicate symptoms preserving order
+    seen = set()
+    unique_indices = []
+    for i, sym in enumerate(result["symptoms"]):
+        sym_lower = sym.lower().strip()
+        if sym_lower not in seen:
+            seen.add(sym_lower)
+            unique_indices.append(i)
+
+    result["symptoms"] = [result["symptoms"][i] for i in unique_indices]
+    result["duration"] = [
+        result["duration"][i] if i < len(result["duration"]) else "unspecified"
+        for i in unique_indices
+    ]
+    result["severity"] = [
+        result["severity"][i] if i < len(result["severity"]) else "unspecified"
+        for i in unique_indices
+    ]
+
+    # Filter out non-medical symptom descriptions
+    valid_indices = [
+        i for i, sym in enumerate(result["symptoms"])
+        if is_valid_symptom(sym)
+    ]
+    # Keep at least one symptom even if filter removes everything
+    if not valid_indices and result["symptoms"]:
+        valid_indices = [0]
+
+    result["symptoms"] = [result["symptoms"][i] for i in valid_indices]
+    result["duration"] = [result["duration"][i] for i in valid_indices]
+    result["severity"] = [result["severity"][i] for i in valid_indices]
+
+    # Pad arrays to same length
+    n = len(result["symptoms"]) or 1
     for key in ("symptoms", "duration", "severity"):
         while len(result[key]) < n:
             result[key].append("unspecified")
+
     return result
 
+
 def severity_badge(s: str) -> str:
     s = (s or "").lower().strip()
-    if not s or s in ("unspecified", "unknown", "n/a", "none", "not mentioned", "—"):
+    if not s or s in ("unspecified", "unknown", "n/a", "none", "not mentioned", "—", "not stated"):
         return '<span style="color:#d1d5db;font-style:italic">—</span>'
-    if any(w in s for w in ("severe", "critical", "extreme", "crushing", "sudden", "acute", "high")):
+    if any(w in s for w in ("severe", "critical", "extreme", "crushing", "sudden", "acute", "high", "sharp")):
         return f'<span style="background:#fef2f2;color:#dc2626;font-weight:600;padding:2px 8px;border-radius:4px;font-size:0.85rem">▲ {s}</span>'
-    if any(w in s for w in ("moderate", "significant", "worsening", "progressive")):
+    if any(w in s for w in ("moderate", "significant", "worsening", "progressive", "persistent")):
         return f'<span style="background:#fffbeb;color:#d97706;font-weight:600;padding:2px 8px;border-radius:4px;font-size:0.85rem">● {s}</span>'
-    if any(w in s for w in ("mild", "slight", "minor", "low", "minimal")):
+    if any(w in s for w in ("mild", "slight", "minor", "low", "minimal", "light")):
         return f'<span style="background:#f0fdf4;color:#16a34a;font-weight:600;padding:2px 8px;border-radius:4px;font-size:0.85rem">▼ {s}</span>'
     return f'<span style="background:#f1f5f9;color:#475569;padding:2px 8px;border-radius:4px;font-size:0.85rem">{s}</span>'
 
-def format_results(result: dict):
+
+def format_duration(d: str) -> str:
+    d = (d or "").lower().strip()
+    if not d or d in ("unspecified", "unknown", "n/a", "none", "not mentioned", "not stated"):
+        return '<span style="color:#d1d5db;font-style:italic">—</span>'
+    return f'<span style="color:#4b5563">{d}</span>'
+
+
+def format_results(result: dict) -> str:
     symptoms = result["symptoms"]
     durations = result["duration"]
     severities = result["severity"]
@@ -120,17 +200,11 @@ def format_results(result: dict):
         sev = severities[i] if i < len(severities) else "unspecified"
         bg = "#fff7f7" if urgent else ("#f8faff" if i % 2 == 0 else "white")
 
-        dur_html = (
-            '<span style="color:#d1d5db;font-style:italic">—</span>'
-            if dur.lower() in ("unspecified", "unknown", "", "n/a")
-            else f'<span style="color:#4b5563">{dur}</span>'
-        )
-
         rows += f"""
         <tr style="background:{bg}">
             <td style="padding:10px 14px;font-weight:600;color:#111827;
                 border-left:3px solid {accent}">{sym}</td>
-            <td style="padding:10px 14px">{dur_html}</td>
+            <td style="padding:10px 14px">{format_duration(dur)}</td>
             <td style="padding:10px 14px">{severity_badge(sev)}</td>
         </tr>"""
 
@@ -157,9 +231,15 @@ def format_results(result: dict):
     </div>
     """
 
+
 def extract(clinical_note: str, state: dict):
     if not clinical_note.strip():
-        return "<p style='color:#9ca3af'>Enter a clinical note to see results.</p>", "{}", state
+        return (
+            "<p style='color:#9ca3af'>Enter a clinical note to see results.</p>",
+            "{}",
+            state,
+            gr.update(visible=True),  # keep warning visible
+        )
 
     load_model()
     prompt = build_prompt(clinical_note)
@@ -183,10 +263,10 @@ def extract(clinical_note: str, state: dict):
         table_html = format_results(result)
         json_out = json.dumps(result, indent=2)
         new_state = {"table": table_html, "json": json_out}
-        return table_html, json_out, new_state
+        return table_html, json_out, new_state, gr.update(visible=False)  # hide warning
     except (json.JSONDecodeError, KeyError, IndexError):
-        err = f"<pre style='color:#dc2626'>Parse error:\n{generated}</pre>"
-        return err, "{}", state
+        err = f"<pre style='color:#dc2626'>Parse error. Raw output:\n{generated}</pre>"
+        return err, "{}", state, gr.update(visible=False)
 
 
 with gr.Blocks(css=CSS, title="ClinicalDistill") as demo:
@@ -194,27 +274,49 @@ with gr.Blocks(css=CSS, title="ClinicalDistill") as demo:
     result_state = gr.State(value={"table": "", "json": "{}"})
 
     gr.HTML("""
-    <h1 id="title" style="font-size:2rem;font-weight:800;background:linear-gradient(135deg,#667eea,#764ba2);
-        -webkit-background-clip:text;-webkit-text-fill-color:transparent;margin-top:1rem">
+    <h1 id="title" style="font-size:2rem;font-weight:800;
+        background:linear-gradient(135deg,#667eea,#764ba2);
+        -webkit-background-clip:text;-webkit-text-fill-color:transparent;
+        margin-top:1rem">
         🏥 ClinicalDistill
     </h1>
-    <p id="subtitle">Structured symptom extraction from clinical notes · Gemma-3-1B fine-tuned with LoRA</p>
+    <p id="subtitle">
+        Structured symptom extraction from clinical notes ·
+        Gemma-3-1B fine-tuned with LoRA
+    </p>
     <div id="stats-row">
-        <div class="stat-card"><div class="stat-val">0.781</div><div class="stat-lbl">F1 Score</div></div>
-        <div class="stat-card"><div class="stat-val">85.7%</div><div class="stat-lbl">Urgent Accuracy</div></div>
-        <div class="stat-card"><div class="stat-val">100%</div><div class="stat-lbl">Valid JSON</div></div>
-        <div class="stat-card"><div class="stat-val">1B</div><div class="stat-lbl">Parameters</div></div>
+        <div class="stat-card">
+            <div class="stat-val">0.781</div>
+            <div class="stat-lbl">F1 Score</div>
+        </div>
+        <div class="stat-card">
+            <div class="stat-val">85.7%</div>
+            <div class="stat-lbl">Urgent Accuracy</div>
+        </div>
+        <div class="stat-card">
+            <div class="stat-val">100%</div>
+            <div class="stat-lbl">Valid JSON</div>
+        </div>
+        <div class="stat-card">
+            <div class="stat-val">1B</div>
+            <div class="stat-lbl">Parameters</div>
+        </div>
     </div>
     """)
 
-    gr.HTML("""
-    <div style="text-align:center;margin-bottom:1rem;padding:0.6rem 1rem;
-        background:#fffbeb;border:1px solid #fde68a;border-radius:8px;
-        color:#92400e;font-size:0.85rem;max-width:600px;margin-left:auto;margin-right:auto">
-        ⏳ Running on CPU — inference takes ~60 seconds. 
-        Results persist after completion.
-    </div>
-    """)
+    # Warning banner — hidden after first inference
+    cpu_warning = gr.HTML(
+        value="""
+        <div style="text-align:center;margin-bottom:1rem;padding:0.6rem 1rem;
+            background:#fffbeb;border:1px solid #fde68a;border-radius:8px;
+            color:#92400e;font-size:0.85rem;max-width:600px;
+            margin-left:auto;margin-right:auto">
+            ⏳ Running on CPU — inference takes ~60 seconds.
+            Results persist after completion.
+        </div>
+        """,
+        visible=True,
+    )
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -226,13 +328,18 @@ with gr.Blocks(css=CSS, title="ClinicalDistill") as demo:
             submit_btn = gr.Button(
                 "⚡ Extract Symptoms",
                 variant="primary",
-                elem_id="submit-btn"
+                elem_id="submit-btn",
+            )
+            gr.Examples(
+                examples=EXAMPLES,
+                inputs=note_input,
+                label="Try an Example",
             )
-            gr.Examples(examples=EXAMPLES, inputs=note_input, label="Try an Example")
 
         with gr.Column(scale=1):
             table_output = gr.HTML(
-                value="<p style='color:#9ca3af;text-align:center;margin-top:2rem'>Results will appear here.</p>",
+                value="<p style='color:#9ca3af;text-align:center;margin-top:2rem'>"
+                      "Results will appear here.</p>",
             )
             with gr.Accordion("Raw JSON Output", open=False):
                 json_output = gr.Code(language="json", label="")
@@ -240,19 +347,21 @@ with gr.Blocks(css=CSS, title="ClinicalDistill") as demo:
     submit_btn.click(
         fn=extract,
         inputs=[note_input, result_state],
-        outputs=[table_output, json_output, result_state]
+        outputs=[table_output, json_output, result_state, cpu_warning],
     )
     note_input.submit(
         fn=extract,
         inputs=[note_input, result_state],
-        outputs=[table_output, json_output, result_state]
+        outputs=[table_output, json_output, result_state, cpu_warning],
     )
 
     gr.HTML("""
     <div style="text-align:center;margin-top:2rem;padding-top:1rem;
         border-top:1px solid #e5e7eb;color:#9ca3af;font-size:0.85rem">
-        ClinicalDistill · Fine-tuned on 145 synthetic clinical examples ·
-        <a href="https://github.com/JanushiShastri/ClinicalDistill" style="color:#667eea">GitHub</a>
+        ClinicalDistill · Fine-tuned on 145 synthetic clinical examples
+        (cardiac, respiratory, neurological, GI) ·
+        <a href="https://github.com/JanushiShastri/ClinicalDistill"
+           style="color:#667eea">GitHub</a>
     </div>
     """)