Upload 2 files

app.py

@@ -38,18 +38,9 @@ CSS = """
 
 HERO = """
 <div id="hero">
-  <h1>🩺 Chart Whisperer</h1>
-  <p>Turn a raw clinical recording into a structured SOAP note. MedASR
-  transcribes the conversation, MedGemma 4B drafts the note with speaker
-  attribution — nothing is fabricated beyond what was actually said.</p>
-  <div class="pills" style="margin-top: 10px;">
-    <span class="pill">🎙️ MedASR</span>
-    <span class="pill">🤖 MedGemma 4B</span>
-    <span class="pill">⚡ ZeroGPU</span>
-  </div>
-  <p style="margin-top:14px; font-size:13px; opacity:.85;">Not a diagnostic
-  tool — for demonstration/research use only. Don't upload real patient
-  data.</p>
+  <h1>Medical Notes Using Medgemma</h1>
+  <p>Turn a raw clinical recording into a structured SOAP note — nothing is
+  fabricated beyond what was actually said.</p>
 </div>
 """
 
@@ -84,7 +75,7 @@ def clear_audio():
     return None, PLACEHOLDER_TRANSCRIPT, PLACEHOLDER_SOAP, PLACEHOLDER_TIMING
 
 
-with gr.Blocks(title="Chart Whisperer") as demo:
+with gr.Blocks(title="Medical Notes Using Medgemma") as demo:
     gr.HTML(HERO)
 
     with gr.Row():
@@ -138,4 +129,4 @@ with gr.Blocks(title="Chart Whisperer") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(theme=THEME, css=CSS)
+    demo.launch(theme=THEME, css=CSS, footer_links=[])

llm.py

@@ -1,6 +1,7 @@
 """Audio -> MedASR transcript -> MedGemma 4B SOAP note pipeline."""
 
 import os
+import re
 import time
 from dataclasses import dataclass
 
@@ -24,8 +25,40 @@ SYSTEM_PROMPT = (
 
 # Forced prefix for the assistant turn: with `continue_final_message=True`,
 # generation resumes mid-turn from this exact text, so there is no token
-# position left for a preamble or transcript restatement to occupy.
-SOAP_PREFIX = "SOAP Note:\n\n**S — Subjective:**"
+# position left for a preamble or transcript restatement to occupy. Only
+# the first header is guaranteed this way — the other three are generated
+# freely and get normalized to match by _normalize_soap_note below.
+SOAP_PREFIX = "S — Subjective:"
+
+_SOAP_HEADERS = [
+    ("S", "Subjective"),
+    ("O", "Objective"),
+    ("A", "Assessment"),
+    ("P", "Plan"),
+]
+_SOAP_TITLE_RE = re.compile(r"(?im)^[ \t]*\**[ \t]*SOAP Note[ \t]*:?\**[ \t]*\n+")
+_SOAP_HEADER_RES = [
+    (
+        re.compile(
+            rf"(?im)^[ \t]*\**[ \t]*(?:{letter}[ \t]*[-—][ \t]*)?{word}[ \t]*:\**"
+        ),
+        f"{letter} — {word}:",
+    )
+    for letter, word in _SOAP_HEADERS
+]
+
+
+def _normalize_soap_note(text: str) -> str:
+    """Force all four section headers to the same 'X — Word:' shape.
+
+    Only the first header is pinned via the forced assistant prefix; the
+    model is free to drift on the rest (e.g. writing 'Plan:' instead of
+    'P — Plan:'), so headers are normalized here rather than trusted.
+    """
+    text = _SOAP_TITLE_RE.sub("", text.strip())
+    for pattern, canonical in _SOAP_HEADER_RES:
+        text = pattern.sub(canonical, text)
+    return text.strip()
 
 _asr_pipe = None
 _llm_model = None
@@ -101,7 +134,7 @@ def generate_soap_note(transcript: str) -> str:
 
     new_tokens = generated_ids[0][input_len:]
     completion = processor.decode(new_tokens, skip_special_tokens=True)
-    return (SOAP_PREFIX + completion).strip()
+    return _normalize_soap_note(SOAP_PREFIX + completion)
 
 
 @dataclass