import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForCausalLM import json import re MODEL_ID = "Janushi/ClinicalDistill-Gemma-1B" tokenizer = None model = None 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"], ["my back's been killing me since last week, hurts way more when i sit, not sure if i pulled something"], ["crushing chest pain radiating to jaw for 30 mins, sweating, feels like something is very wrong"], ["kid woke up hot last night, been sneezing a lot, seems okay otherwise just a little cranky"], ["stomach's been acting up since yesterday, went to the bathroom like 4 times, feeling really drained"], ["these headaches keep coming back, nothing crazy but annoying, sometimes feel dizzy too"], ] CSS = """ #title { text-align: center; margin-bottom: 0.5rem; } #subtitle { text-align: center; color: #6b7280; margin-bottom: 1.5rem; font-size: 0.95rem; } #stats-row { display: flex; gap: 1rem; justify-content: center; margin-bottom: 1.5rem; flex-wrap: wrap; } .stat-card { background: linear-gradient(135deg, #667eea20, #764ba220); border: 1px solid #667eea40; border-radius: 12px; padding: 0.6rem 1.2rem; text-align: center; min-width: 110px; } .stat-val { font-size: 1.4rem; font-weight: 700; color: #4f46e5; } .stat-lbl { font-size: 0.75rem; color: #6b7280; } #submit-btn { background: linear-gradient(135deg, #667eea, #764ba2) !important; border: none !important; font-size: 1rem !important; padding: 0.75rem !important; } #urgent-badge-yes { background: #fef2f2; border: 1px solid #fca5a5; color: #dc2626; border-radius: 8px; padding: 0.5rem 1rem; font-weight: 600; text-align: center; margin-top: 0.5rem; } #urgent-badge-no { background: #f0fdf4; border: 1px solid #86efac; color: #16a34a; border-radius: 8px; padding: 0.5rem 1rem; font-weight: 600; text-align: center; margin-top: 0.5rem; } 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: return tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float32, # float32 for CPU device_map="cpu", ) model.eval() def build_prompt(clinical_note: str) -> str: return ( f"\n{INSTRUCTION}\n\n\n" f"\n{clinical_note.strip()}\n\n\n" f"\n" ) def parse_output(raw: str) -> dict: raw = raw.split("")[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 # 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", "—", "not stated"): return '—' if any(w in s for w in ("severe", "critical", "extreme", "crushing", "sudden", "acute", "high", "sharp")): return f'▲ {s}' if any(w in s for w in ("moderate", "significant", "worsening", "progressive", "persistent")): return f'● {s}' if any(w in s for w in ("mild", "slight", "minor", "low", "minimal", "light")): return f'▼ {s}' return f'{s}' 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 '—' return f'{d}' def format_results(result: dict) -> str: symptoms = result["symptoms"] durations = result["duration"] severities = result["severity"] urgent = result.get("urgent", False) accent = "#dc2626" if urgent else "#667eea" rows = "" for i, sym in enumerate(symptoms): dur = durations[i] if i < len(durations) else "unspecified" sev = severities[i] if i < len(severities) else "unspecified" bg = "#fff7f7" if urgent else ("#f8faff" if i % 2 == 0 else "white") rows += f""" {sym} {format_duration(dur)} {severity_badge(sev)} """ urgent_html = ( '
🚨 URGENT — Immediate attention may be required
' if urgent else '
✅ NON-URGENT — Routine follow-up appropriate
' ) return f"""
{rows}
Symptom Duration Severity
{urgent_html}
""" def extract(clinical_note: str, state: dict): if not clinical_note.strip(): return ( "

Enter a clinical note to see results.

", "{}", state, gr.update(visible=True), # keep warning visible ) load_model() prompt = build_prompt(clinical_note) inputs = tokenizer(prompt, return_tensors="pt") with torch.inference_mode(): output_ids = model.generate( **inputs, max_new_tokens=256, do_sample=False, pad_token_id=tokenizer.eos_token_id, ) generated = tokenizer.decode( output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True, ) try: result = parse_output(generated) 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, gr.update(visible=False) # hide warning except (json.JSONDecodeError, KeyError, IndexError): err = f"
Parse error. Raw output:\n{generated}
" return err, "{}", state, gr.update(visible=False) with gr.Blocks(css=CSS, title="ClinicalDistill") as demo: result_state = gr.State(value={"table": "", "json": "{}"}) gr.HTML("""

🏥 ClinicalDistill

Structured symptom extraction from clinical notes · Gemma-3-1B fine-tuned with LoRA

0.781
F1 Score
85.7%
Urgent Accuracy
100%
Valid JSON
1B
Parameters
""") # Warning banner — hidden after first inference cpu_warning = gr.HTML( value="""
⏳ Running on CPU — inference takes ~60 seconds. Results persist after completion.
""", visible=True, ) with gr.Row(): with gr.Column(scale=1): note_input = gr.Textbox( label="Clinical Note", placeholder="e.g. been feeling off, chest feels weird and i get tired walking around...", lines=7, ) submit_btn = gr.Button( "⚡ Extract Symptoms", variant="primary", elem_id="submit-btn", ) gr.Examples( examples=EXAMPLES, inputs=note_input, label="Try an Example", ) with gr.Column(scale=1): table_output = gr.HTML( value="

" "Results will appear here.

", ) with gr.Accordion("Raw JSON Output", open=False): json_output = gr.Code(language="json", label="") submit_btn.click( fn=extract, inputs=[note_input, 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, cpu_warning], ) gr.HTML("""
ClinicalDistill · Fine-tuned on 145 synthetic clinical examples (cardiac, respiratory, neurological, GI) · GitHub
""") demo.launch(ssr_mode=False)