sugatobagchi commited on
Commit
ecfd585
·
verified ·
1 Parent(s): a4add95

Upload 8 files

Browse files
.gitattributes CHANGED
@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ test_audio/clips/D0420-S1-T01_clip.wav filter=lfs diff=lfs merge=lfs -text
37
+ test_audio/clips/D0420-S2-T01_clip.wav filter=lfs diff=lfs merge=lfs -text
38
+ test_audio/clips/D0420-S3-T01_clip.wav filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -1,15 +1,41 @@
1
- ---
2
- title: Medical Notes Using Medgemma
3
- emoji: 🏢
4
- colorFrom: pink
5
- colorTo: yellow
6
- sdk: gradio
7
- sdk_version: 6.20.0
8
- python_version: '3.12'
9
- app_file: app.py
10
- pinned: false
11
- license: apache-2.0
12
- short_description: Convert Medical Conversations into Medical Notes.
13
- ---
14
-
15
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # AI Medical Scribe
2
+
3
+ Pipeline: audio file -> MedASR (transcription) -> MedGemma 4B, `google/medgemma-4b-it`
4
+ (SOAP note) -> Gradio UI.
5
+
6
+ Both models are gated on the Hugging Face Hub, so whichever environment runs
7
+ this needs an HF token with access to `google/medasr` and `google/medgemma-4b-it`.
8
+
9
+ ## Setup (local)
10
+
11
+ ```bash
12
+ python3.12 -m venv venv
13
+ source venv/bin/activate
14
+ pip install -r requirements.txt
15
+ hf auth login # token needs access to the gated google/medasr and google/medgemma-4b-it repos
16
+ ```
17
+
18
+ Model weights are pulled straight from the Hub on first run (no local
19
+ `models/` download step needed anymore). To point at a local copy instead,
20
+ set `MEDASR_MODEL_ID` / `MEDGEMMA_MODEL_ID` to a local directory.
21
+
22
+ ## Run (local)
23
+
24
+ ```bash
25
+ python app.py
26
+ ```
27
+
28
+ Opens a local Gradio app: pick a wav/mp3 file, click Run, see the MedASR
29
+ transcript, the generated SOAP note, and per-step timing.
30
+
31
+ ## Deploying to Hugging Face Spaces (ZeroGPU)
32
+
33
+ 1. Create a Space with SDK "Gradio" and hardware "ZeroGPU".
34
+ 2. Push this repo's contents to the Space: `app.py`, `llm.py`,
35
+ `requirements.txt`, `README.md`, `test_audio/clips/*.wav`.
36
+ 3. In the Space's Settings -> Variables and secrets, add a secret `HF_TOKEN`
37
+ set to a token with access to the gated `google/medasr` and
38
+ `google/medgemma-4b-it` repos. `huggingface_hub` reads this automatically.
39
+ 4. `llm.py`'s `run_pipeline` is decorated with `@spaces.GPU`, so ZeroGPU
40
+ attaches a GPU only for the duration of that call; both the ASR and
41
+ SOAP-note models load lazily on first invocation.
app.py ADDED
@@ -0,0 +1,499 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """AI medical scribe: audio -> MedASR transcript -> MedGemma SOAP note."""
2
+
3
+ import os
4
+
5
+ import gradio as gr
6
+
7
+ from llm import run_pipeline
8
+
9
+ _CLIPS_DIR = os.path.join(
10
+ os.path.dirname(os.path.abspath(__file__)), "test_audio", "clips"
11
+ )
12
+
13
+ SAMPLES = {
14
+ "Sample 1 — Psychiatric referral": os.path.join(
15
+ _CLIPS_DIR, "D0420-S1-T01_clip.wav"
16
+ ),
17
+ "Sample 2 — Mood check-in": os.path.join(_CLIPS_DIR, "D0420-S2-T01_clip.wav"),
18
+ "Sample 3 — Sleep & nightmares": os.path.join(_CLIPS_DIR, "D0420-S3-T01_clip.wav"),
19
+ }
20
+
21
+ # Hardcoded light palette only — no Gradio theme / dark mode.
22
+ CUSTOM_CSS = """
23
+ :root, :root.dark, :root .dark, .dark, html, body {
24
+ color-scheme: light !important;
25
+
26
+ --radius: 0.9rem;
27
+ --bg: #f6f7fb;
28
+ --bg-glow: radial-gradient(ellipse 70% 40% at 15% -10%, rgba(59,130,246,0.07), transparent),
29
+ radial-gradient(ellipse 60% 40% at 100% 0%, rgba(139,92,246,0.06), transparent);
30
+ --card: #ffffff;
31
+ --card-alt: #f2f4f9;
32
+ --border: #e3e7f0;
33
+ --fg: #10131c;
34
+ --muted: #667085;
35
+ --primary: #3b82f6;
36
+ --primary-hover: #2563eb;
37
+ --shadow: 0 1px 2px rgba(16,24,40,0.06), 0 1px 0 rgba(16,24,40,0.02);
38
+
39
+ /* Override Gradio tokens so dark class / system preference cannot darken the UI */
40
+ --body-background-fill: #f6f7fb !important;
41
+ --background-fill-primary: #f6f7fb !important;
42
+ --background-fill-secondary: #ffffff !important;
43
+ --block-background-fill: #ffffff !important;
44
+ --block-border-color: #e3e7f0 !important;
45
+ --border-color-primary: #e3e7f0 !important;
46
+ --body-text-color: #10131c !important;
47
+ --body-text-color-subdued: #667085 !important;
48
+ --neutral-50: #f8fafc !important;
49
+ --neutral-100: #f1f5f9 !important;
50
+ --neutral-200: #e2e8f0 !important;
51
+ --neutral-300: #cbd5e1 !important;
52
+ --neutral-400: #94a3b8 !important;
53
+ --neutral-500: #64748b !important;
54
+ --neutral-600: #475569 !important;
55
+ --neutral-700: #334155 !important;
56
+ --neutral-800: #1e293b !important;
57
+ --neutral-900: #0f172a !important;
58
+ --neutral-950: #020617 !important;
59
+ --input-background-fill: #f2f4f9 !important;
60
+ --input-border-color: #e3e7f0 !important;
61
+ }
62
+
63
+ * { box-sizing: border-box; }
64
+
65
+ html, body, gradio-app, .gradio-container, .main, .wrap, .contain,
66
+ #root, .app, .fillable, .dark, .dark body, body.dark {
67
+ background: #f6f7fb !important;
68
+ background-color: #f6f7fb !important;
69
+ background-image: var(--bg-glow) !important;
70
+ background-attachment: fixed !important;
71
+ color: #10131c !important;
72
+ font-family: 'Inter', system-ui, sans-serif !important;
73
+ }
74
+
75
+ footer { display: none !important; }
76
+ .gradio-container {
77
+ max-width: 100% !important;
78
+ width: 100% !important;
79
+ margin: 0 !important;
80
+ padding: 0 clamp(1.25rem, 5vw, 4rem) !important;
81
+ min-height: 100vh !important;
82
+ }
83
+
84
+ /* Strip Gradio block chrome; only .card should paint a background. */
85
+ .block, .form, .gr-group, div[data-testid="block"], .gradio-container .padded {
86
+ background: transparent !important;
87
+ border: none !important;
88
+ box-shadow: none !important;
89
+ }
90
+
91
+ #top-bar {
92
+ display: flex;
93
+ align-items: center;
94
+ justify-content: space-between;
95
+ padding: 1.5rem 0 0 0;
96
+ }
97
+ #brand {
98
+ display: flex;
99
+ align-items: center;
100
+ gap: 0.6rem;
101
+ font-family: 'Outfit', sans-serif;
102
+ font-weight: 700;
103
+ font-size: 1.05rem;
104
+ color: #10131c !important;
105
+ }
106
+ #brand .logo-dot {
107
+ width: 10px; height: 10px; border-radius: 3px;
108
+ background: linear-gradient(135deg, #3b82f6, #8b5cf6);
109
+ box-shadow: 0 0 10px rgba(59,130,246,0.6);
110
+ }
111
+
112
+ #app-header {
113
+ text-align: center;
114
+ padding: 2.75rem 1rem 2.25rem 1rem;
115
+ }
116
+ #app-header .eyebrow {
117
+ display: inline-block;
118
+ font-size: 0.72rem;
119
+ font-weight: 600;
120
+ letter-spacing: 0.12em;
121
+ text-transform: uppercase;
122
+ color: #2563eb !important;
123
+ background: rgba(59,130,246,0.12);
124
+ border: 1px solid rgba(59,130,246,0.3);
125
+ padding: 0.3rem 0.9rem;
126
+ border-radius: 999px;
127
+ margin-bottom: 1.1rem;
128
+ }
129
+ #app-header h1 {
130
+ font-family: 'Outfit', sans-serif;
131
+ font-weight: 800;
132
+ font-size: clamp(2rem, 4vw, 2.9rem);
133
+ letter-spacing: -0.02em;
134
+ margin: 0 0 0.6rem 0;
135
+ color: #10131c !important;
136
+ }
137
+ #app-header p {
138
+ color: #667085 !important;
139
+ font-size: 1.02rem;
140
+ max-width: 640px;
141
+ margin: 0 auto;
142
+ line-height: 1.6;
143
+ }
144
+
145
+ .card {
146
+ background: #ffffff !important;
147
+ border: 1px solid #e3e7f0 !important;
148
+ border-radius: var(--radius) !important;
149
+ box-shadow: var(--shadow);
150
+ }
151
+
152
+ #control-row { align-items: stretch; gap: 1rem; }
153
+
154
+ #sample-panel { padding: 1.25rem !important; margin-bottom: 1rem; }
155
+ #sample-radio .wrap {
156
+ display: flex !important;
157
+ flex-wrap: wrap;
158
+ justify-content: center;
159
+ gap: 0.75rem;
160
+ background: transparent !important;
161
+ border: none !important;
162
+ }
163
+ #sample-radio label {
164
+ position: relative;
165
+ background: #f9fafc !important;
166
+ border: 1.5px solid #e3e7f0 !important;
167
+ border-radius: 0.65rem !important;
168
+ padding: 0.7rem 1.1rem 0.7rem 2.15rem !important;
169
+ font-size: 0.85rem !important;
170
+ font-weight: 600 !important;
171
+ text-transform: none !important;
172
+ letter-spacing: normal !important;
173
+ color: #10131c !important;
174
+ box-shadow: 0 1px 2px rgba(16,24,40,0.05);
175
+ cursor: pointer;
176
+ transition: background 0.15s ease, border-color 0.15s ease, box-shadow 0.15s ease, transform 0.1s ease;
177
+ }
178
+ /* custom radio dot, drawn in place of the hidden native input */
179
+ #sample-radio label::before {
180
+ content: "";
181
+ position: absolute;
182
+ left: 0.85rem;
183
+ top: 50%;
184
+ transform: translateY(-50%);
185
+ width: 15px;
186
+ height: 15px;
187
+ border-radius: 50%;
188
+ border: 1.5px solid #cbd5e1;
189
+ background: #ffffff;
190
+ transition: border-color 0.15s ease, background 0.15s ease, box-shadow 0.15s ease;
191
+ }
192
+ #sample-radio label:hover {
193
+ border-color: #93c5fd !important;
194
+ box-shadow: 0 2px 8px rgba(59,130,246,0.12);
195
+ transform: translateY(-1px);
196
+ }
197
+ #sample-radio label:hover::before { border-color: #3b82f6; }
198
+ #sample-radio label:has(input:checked) {
199
+ background: #eaf1ff !important;
200
+ border-color: #3b82f6 !important;
201
+ color: #1d4ed8 !important;
202
+ box-shadow: 0 3px 12px rgba(59,130,246,0.25);
203
+ }
204
+ #sample-radio label:has(input:checked)::before {
205
+ border-color: #3b82f6;
206
+ background: #3b82f6;
207
+ box-shadow: inset 0 0 0 3px #ffffff;
208
+ }
209
+ #sample-radio input[type="radio"] { display: none; }
210
+
211
+ label span, .gr-form label {
212
+ color: #667085 !important;
213
+ font-weight: 500 !important;
214
+ font-size: 0.78rem !important;
215
+ letter-spacing: 0.03em;
216
+ text-transform: uppercase;
217
+ }
218
+
219
+ #audio-upload {
220
+ background: #ffffff !important;
221
+ border: 1px dashed #e3e7f0 !important;
222
+ border-radius: var(--radius) !important;
223
+ color: #10131c !important;
224
+ min-height: 220px;
225
+ }
226
+ #audio-upload:hover { border-color: #3b82f6 !important; }
227
+ #audio-upload .audio-container,
228
+ #audio-upload button.svelte-8prmba {
229
+ width: 100% !important;
230
+ background: transparent !important;
231
+ }
232
+ #audio-upload .wrap.svelte-1vmd51o {
233
+ width: 100% !important;
234
+ height: 100% !important;
235
+ background: transparent !important;
236
+ color: #10131c !important;
237
+ gap: 0.35rem;
238
+ }
239
+ #audio-upload .wrap.svelte-1vmd51o .icon-wrap {
240
+ color: #3b82f6 !important;
241
+ width: 34px !important;
242
+ height: 34px !important;
243
+ margin-bottom: 0.25rem;
244
+ }
245
+ #audio-upload .wrap.svelte-1vmd51o .or {
246
+ color: #94a3b8 !important;
247
+ font-size: 0.85rem !important;
248
+ }
249
+
250
+ button.primary, #run-btn button {
251
+ background: linear-gradient(180deg, #2563eb, #3b82f6) !important;
252
+ border: none !important;
253
+ border-radius: 0.75rem !important;
254
+ color: white !important;
255
+ font-weight: 700 !important;
256
+ font-size: 1rem !important;
257
+ letter-spacing: 0.01em;
258
+ box-shadow: 0 6px 20px rgba(59,130,246,0.35);
259
+ transition: transform 0.15s ease, box-shadow 0.15s ease;
260
+ }
261
+ button.primary:hover, #run-btn button:hover {
262
+ transform: translateY(-1px);
263
+ box-shadow: 0 8px 26px rgba(59,130,246,0.5);
264
+ }
265
+ #run-btn button { width: 100%; height: 3.1rem !important; }
266
+ #run-row { margin-top: 1rem; }
267
+
268
+ textarea, input[type="text"] {
269
+ background: #f2f4f9 !important;
270
+ border: 1px solid #e3e7f0 !important;
271
+ border-radius: 0.6rem !important;
272
+ color: #10131c !important;
273
+ font-family: 'Inter', monospace !important;
274
+ font-size: 0.92rem !important;
275
+ line-height: 1.55 !important;
276
+ }
277
+
278
+ #transcript-panel, #soap-panel, #timing-panel {
279
+ padding: 1.25rem !important;
280
+ }
281
+
282
+ .panel-title {
283
+ font-family: 'Outfit', sans-serif;
284
+ font-weight: 600;
285
+ font-size: 1.02rem;
286
+ color: #10131c !important;
287
+ margin-bottom: 0.9rem;
288
+ display: flex;
289
+ align-items: center;
290
+ gap: 0.5rem;
291
+ }
292
+ .panel-title .dot {
293
+ width: 8px; height: 8px; border-radius: 50%;
294
+ }
295
+ .dot-blue { background: #3b82f6; box-shadow: 0 0 8px #3b82f6; }
296
+ .dot-violet { background: #8b5cf6; box-shadow: 0 0 8px #8b5cf6; }
297
+ .dot-emerald { background: #10b981; box-shadow: 0 0 8px #10b981; }
298
+
299
+ #soap-panel .prose, #soap-panel p, #soap-panel li, #soap-panel strong {
300
+ color: #10131c !important;
301
+ }
302
+ #soap-panel { min-height: 340px; }
303
+
304
+ .timing-badges { display: flex; gap: 0.75rem; flex-wrap: wrap; }
305
+ .timing-badge {
306
+ flex: 1;
307
+ min-width: 150px;
308
+ background: #f2f4f9;
309
+ border: 1px solid #e3e7f0;
310
+ border-radius: 0.75rem;
311
+ padding: 0.85rem 1rem;
312
+ }
313
+ .timing-badge .label {
314
+ font-size: 0.72rem;
315
+ text-transform: uppercase;
316
+ letter-spacing: 0.08em;
317
+ color: #667085;
318
+ margin-bottom: 0.3rem;
319
+ }
320
+ .timing-badge .value {
321
+ font-family: 'Outfit', sans-serif;
322
+ font-size: 1.4rem;
323
+ font-weight: 600;
324
+ color: #10131c;
325
+ }
326
+ .timing-badge .value span { font-size: 0.95rem; color: #667085; font-weight: 500; }
327
+ """
328
+
329
+ HEAD = """
330
+ <link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=Outfit:wght@500;600;700;800&display=swap">
331
+ <script>
332
+ (() => {
333
+ const LIGHT = "#f6f7fb";
334
+ let painting = false;
335
+ const paint = () => {
336
+ if (painting) return;
337
+ painting = true;
338
+ try {
339
+ document.documentElement.classList.remove("dark");
340
+ document.body?.classList.remove("dark");
341
+ document.querySelectorAll(".dark").forEach((el) => el.classList.remove("dark"));
342
+ for (const el of [document.documentElement, document.body, document.querySelector("gradio-app")]) {
343
+ if (!el) continue;
344
+ el.style.setProperty("background", LIGHT, "important");
345
+ el.style.setProperty("background-color", LIGHT, "important");
346
+ el.style.setProperty("color-scheme", "light", "important");
347
+ }
348
+ } finally {
349
+ painting = false;
350
+ }
351
+ };
352
+ const url = new URL(window.location.href);
353
+ if (url.searchParams.get("__theme") !== "light") {
354
+ url.searchParams.set("__theme", "light");
355
+ window.history.replaceState({}, "", url);
356
+ }
357
+ paint();
358
+ document.addEventListener("DOMContentLoaded", paint);
359
+ new MutationObserver((mutations) => {
360
+ const needs = mutations.some((m) => {
361
+ const el = m.target;
362
+ return el.classList?.contains("dark") || (el.style && !String(el.style.background || "").includes(LIGHT));
363
+ });
364
+ if (needs) paint();
365
+ }).observe(document.documentElement, {
366
+ attributes: true,
367
+ attributeFilter: ["class", "style"],
368
+ subtree: true,
369
+ });
370
+ })();
371
+ </script>
372
+ """
373
+
374
+ TOPBAR_HTML = """
375
+ <div id="top-bar">
376
+ <div id="brand"><span class="logo-dot"></span>MedScribe</div>
377
+ </div>
378
+ """
379
+
380
+ HEADER_HTML = """
381
+ <div id="app-header">
382
+ <span class="eyebrow">MedASR + MedGemma 4B</span>
383
+ <h1>AI Medical Scribe</h1>
384
+ <p>Turn a raw clinical recording into a structured SOAP note.</p>
385
+ </div>
386
+ """
387
+
388
+
389
+ def timing_html(transcription_s: float, generation_s: float) -> str:
390
+ total = transcription_s + generation_s
391
+ return f"""
392
+ <div class="timing-badges">
393
+ <div class="timing-badge">
394
+ <div class="label">Transcription</div>
395
+ <div class="value">{transcription_s:.2f}<span>s</span></div>
396
+ </div>
397
+ <div class="timing-badge">
398
+ <div class="label">SOAP Generation</div>
399
+ <div class="value">{generation_s:.2f}<span>s</span></div>
400
+ </div>
401
+ <div class="timing-badge">
402
+ <div class="label">Total</div>
403
+ <div class="value">{total:.2f}<span>s</span></div>
404
+ </div>
405
+ </div>
406
+ """
407
+
408
+
409
+ def pick_sample(label):
410
+ return SAMPLES.get(label)
411
+
412
+
413
+ def run(audio_file):
414
+ if audio_file is None:
415
+ return (
416
+ "No audio file selected.",
417
+ "*Select an audio file and click Run.*",
418
+ timing_html(0, 0),
419
+ )
420
+
421
+ result = run_pipeline(audio_file)
422
+ return (
423
+ result.transcript,
424
+ result.soap_note,
425
+ timing_html(result.transcription_seconds, result.generation_seconds),
426
+ )
427
+
428
+
429
+ with gr.Blocks(title="AI Medical Scribe") as demo:
430
+ gr.HTML(TOPBAR_HTML)
431
+ gr.HTML(HEADER_HTML)
432
+
433
+ with gr.Column(elem_classes=["card"], elem_id="sample-panel"):
434
+ gr.HTML(
435
+ '<div class="panel-title"><span class="dot dot-blue"></span>Choose a sample recording</div>'
436
+ )
437
+ sample_picker = gr.Radio(
438
+ choices=list(SAMPLES.keys()),
439
+ value=None,
440
+ show_label=False,
441
+ elem_id="sample-radio",
442
+ )
443
+
444
+ with gr.Column(elem_id="control-row"):
445
+ audio_input = gr.Audio(
446
+ label="Audio recording",
447
+ sources=["upload"],
448
+ type="filepath",
449
+ elem_id="audio-upload",
450
+ )
451
+ with gr.Row(elem_id="run-row"):
452
+ run_button = gr.Button(
453
+ "Run pipeline", variant="primary", elem_id="run-btn", size="lg"
454
+ )
455
+
456
+ with gr.Row():
457
+ with gr.Column(elem_classes=["card"], elem_id="transcript-panel"):
458
+ gr.HTML(
459
+ '<div class="panel-title"><span class="dot dot-blue"></span>MedASR Transcript</div>'
460
+ )
461
+ transcript_output = gr.Textbox(
462
+ show_label=False,
463
+ lines=15,
464
+ container=False,
465
+ placeholder="Transcript will appear here…",
466
+ )
467
+ with gr.Column(elem_classes=["card"], elem_id="soap-panel"):
468
+ gr.HTML(
469
+ '<div class="panel-title"><span class="dot dot-violet"></span>SOAP Note — MedGemma 4B</div>'
470
+ )
471
+ soap_output = gr.Markdown("*SOAP note will appear here…*")
472
+
473
+ with gr.Column(elem_classes=["card"], elem_id="timing-panel"):
474
+ gr.HTML(
475
+ '<div class="panel-title"><span class="dot dot-emerald"></span>Step Timing</div>'
476
+ )
477
+ timing_output = gr.HTML(timing_html(0, 0))
478
+
479
+ # Picking a sample loads it into the player (so it can be listened to
480
+ # immediately); the player's own .change then kicks off the pipeline,
481
+ # so transcription/generation run while playback is already available.
482
+ sample_picker.change(fn=pick_sample, inputs=[sample_picker], outputs=[audio_input])
483
+ audio_input.change(
484
+ fn=run,
485
+ inputs=[audio_input],
486
+ outputs=[transcript_output, soap_output, timing_output],
487
+ )
488
+ run_button.click(
489
+ fn=run,
490
+ inputs=[audio_input],
491
+ outputs=[transcript_output, soap_output, timing_output],
492
+ )
493
+
494
+ if __name__ == "__main__":
495
+ demo.launch(
496
+ css=CUSTOM_CSS,
497
+ head=HEAD,
498
+ footer_links=[],
499
+ )
llm.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Audio -> MedASR transcript -> MedGemma 4B SOAP note pipeline."""
2
+
3
+ import os
4
+ import time
5
+ from dataclasses import dataclass
6
+
7
+ import spaces
8
+ import torch
9
+ from transformers import AutoModelForImageTextToText, AutoProcessor, pipeline
10
+
11
+ ASR_MODEL_ID = os.environ.get("MEDASR_MODEL_ID", "google/medasr")
12
+ LLM_MODEL_ID = os.environ.get("MEDGEMMA_MODEL_ID", "google/medgemma-4b-it")
13
+
14
+ SYSTEM_PROMPT = (
15
+ "You are a clinical documentation engine. Convert the transcript into a "
16
+ "SOAP note (Subjective, Objective, Assessment, Plan). Infer which "
17
+ "statements come from the doctor versus the patient based on context "
18
+ "(questions, clinical observations vs. symptom descriptions). Do not "
19
+ "fabricate any detail not present in the transcript. If speaker "
20
+ "attribution is unclear, mark it as unclear rather than guessing "
21
+ "confidently."
22
+ )
23
+
24
+ _asr_pipe = None
25
+ _llm_model = None
26
+ _llm_processor = None
27
+
28
+
29
+ def _device_and_dtype():
30
+ if torch.cuda.is_available():
31
+ return "cuda", torch.bfloat16
32
+ if torch.backends.mps.is_available():
33
+ return "mps", torch.float16
34
+ return "cpu", torch.float32
35
+
36
+
37
+ def get_asr_pipeline():
38
+ global _asr_pipe
39
+ if _asr_pipe is None:
40
+ device, dtype = _device_and_dtype()
41
+ _asr_pipe = pipeline(
42
+ "automatic-speech-recognition",
43
+ model=ASR_MODEL_ID,
44
+ device=device,
45
+ dtype=dtype,
46
+ )
47
+ return _asr_pipe
48
+
49
+
50
+ def get_llm():
51
+ global _llm_model, _llm_processor
52
+ if _llm_model is None:
53
+ device, dtype = _device_and_dtype()
54
+ _llm_model = AutoModelForImageTextToText.from_pretrained(
55
+ LLM_MODEL_ID, dtype=dtype, device_map=device
56
+ )
57
+ _llm_processor = AutoProcessor.from_pretrained(LLM_MODEL_ID)
58
+ return _llm_model, _llm_processor
59
+
60
+
61
+ def transcribe(audio_path: str) -> str:
62
+ """Transcribe an audio file to text using MedASR, chunked for long audio."""
63
+ pipe = get_asr_pipeline()
64
+ result = pipe(audio_path, chunk_length_s=20, stride_length_s=2)
65
+ return result["text"].strip()
66
+
67
+
68
+ def generate_soap_note(transcript: str) -> str:
69
+ model, processor = get_llm()
70
+ messages = [
71
+ {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
72
+ {"role": "user", "content": [{"type": "text", "text": transcript}]},
73
+ ]
74
+ inputs = processor.apply_chat_template(
75
+ messages,
76
+ add_generation_prompt=True,
77
+ tokenize=True,
78
+ return_dict=True,
79
+ return_tensors="pt",
80
+ ).to(model.device)
81
+
82
+ input_len = inputs["input_ids"].shape[-1]
83
+
84
+ with torch.inference_mode():
85
+ generated_ids = model.generate(
86
+ **inputs,
87
+ max_new_tokens=1024,
88
+ do_sample=False,
89
+ repetition_penalty=1.3,
90
+ )
91
+
92
+ new_tokens = generated_ids[0][input_len:]
93
+ return processor.decode(new_tokens, skip_special_tokens=True).strip()
94
+
95
+
96
+ @dataclass
97
+ class PipelineResult:
98
+ transcript: str
99
+ soap_note: str
100
+ transcription_seconds: float
101
+ generation_seconds: float
102
+
103
+
104
+ @spaces.GPU(duration=120)
105
+ def run_pipeline(audio_path: str) -> PipelineResult:
106
+ t0 = time.perf_counter()
107
+ transcript = transcribe(audio_path)
108
+ t1 = time.perf_counter()
109
+
110
+ soap_note = generate_soap_note(transcript)
111
+ t2 = time.perf_counter()
112
+
113
+ return PipelineResult(
114
+ transcript=transcript,
115
+ soap_note=soap_note,
116
+ transcription_seconds=t1 - t0,
117
+ generation_seconds=t2 - t1,
118
+ )
requirements.txt ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch==2.13.0
2
+ torchaudio==2.11.0
3
+ transformers==5.14.1
4
+ accelerate==1.14.0
5
+ soundfile==0.14.0
6
+ librosa==0.11.0
7
+ huggingface_hub==1.24.0
8
+ gradio==6.20.0
9
+ spaces
10
+ sentencepiece
11
+ pillow
test_audio/.DS_Store ADDED
Binary file (6.15 kB). View file
 
test_audio/clips/D0420-S1-T01_clip.wav ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3aca81a6750e0e091fa7daec5a3046513f8fdded0dcf4b3cbef45ef86b841434
3
+ size 2880154
test_audio/clips/D0420-S2-T01_clip.wav ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0b990a7cf8d577e73ab47cd5580f95ab053feee176fad7dc54b007f1def38b97
3
+ size 2880154
test_audio/clips/D0420-S3-T01_clip.wav ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:163265a5609dd0aec3c322786b5296275166350e708523b2ac9c3e44ae08a797
3
+ size 2880154