app.py

"""AI medical scribe: audio -> MedASR transcript -> MedGemma SOAP note."""

import os

import gradio as gr

from llm import run_pipeline

_CLIPS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_audio", "clips")

SAMPLES = [
    ("Sample 1 — Psychiatric referral", os.path.join(_CLIPS_DIR, "D0420-S1-T01_clip.wav")),
    ("Sample 2 — Mood check-in", os.path.join(_CLIPS_DIR, "D0420-S2-T01_clip.wav")),
    ("Sample 3 — Sleep & nightmares", os.path.join(_CLIPS_DIR, "D0420-S3-T01_clip.wav")),
]

THEME = gr.themes.Soft(primary_hue=gr.themes.colors.blue)

CSS = """
.gradio-container { max-width: 1100px !important; }
#hero {
    background: linear-gradient(135deg, #2563eb 0%, #7c3aed 100%);
    color: #fff;
    border-radius: 16px;
    padding: 28px;
}
#hero h1 { margin: 0 0 8px 0; font-size: 31px; font-weight: 800; }
.pill {
    display: inline-block;
    background: rgba(255,255,255,.2);
    border: 1px solid rgba(255,255,255,.4);
    padding: 5px 13px;
    border-radius: 999px;
    font-size: 12.5px;
    margin: 0 8px 6px 0;
}
"""

HERO = """
<div id="hero">
  <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>
"""


def timing_md(transcription_s: float, generation_s: float) -> str:
    total = transcription_s + generation_s
    return (
        f"**Transcription:** {transcription_s:.2f}s &nbsp;·&nbsp; "
        f"**SOAP generation:** {generation_s:.2f}s &nbsp;·&nbsp; "
        f"**Total:** {total:.2f}s"
    )


PLACEHOLDER_TRANSCRIPT = ""
PLACEHOLDER_SOAP = "*Select an audio file and click Run.*"
PLACEHOLDER_TIMING = ""


def run(audio_file):
    if audio_file is None:
        return "No audio file selected.", PLACEHOLDER_SOAP, PLACEHOLDER_TIMING

    result = run_pipeline(audio_file)
    return (
        result.transcript,
        result.soap_note,
        timing_md(result.transcription_seconds, result.generation_seconds),
    )


def clear_audio():
    return None, PLACEHOLDER_TRANSCRIPT, PLACEHOLDER_SOAP, PLACEHOLDER_TIMING


with gr.Blocks(title="Medical Notes Using Medgemma") as demo:
    gr.HTML(HERO)

    with gr.Row():
        sample_buttons = [gr.Button(label, size="sm") for label, _ in SAMPLES]

    with gr.Row():
        with gr.Column(scale=3):
            audio_input = gr.Audio(
                label="Audio recording",
                sources=["upload"],
                type="filepath",
                buttons=[],
            )
        with gr.Column(scale=1, min_width=160):
            run_button = gr.Button("Run pipeline", variant="primary", size="lg")
            clear_button = gr.Button("Clear audio", size="sm")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### MedASR Transcript")
            with gr.Group():
                transcript_output = gr.Textbox(
                    show_label=False,
                    lines=15,
                    container=False,
                    placeholder="Transcript will appear here…",
                )
        with gr.Column():
            gr.Markdown("### SOAP Note — MedGemma 4B")
            with gr.Group():
                soap_output = gr.Markdown(PLACEHOLDER_SOAP)

    timing_output = gr.Markdown(PLACEHOLDER_TIMING)

    for btn, (_, path) in zip(sample_buttons, SAMPLES):
        btn.click(lambda p=path: p, outputs=[audio_input])

    audio_input.change(
        fn=run,
        inputs=[audio_input],
        outputs=[transcript_output, soap_output, timing_output],
    )
    run_button.click(
        fn=run,
        inputs=[audio_input],
        outputs=[transcript_output, soap_output, timing_output],
    )
    clear_button.click(
        fn=clear_audio,
        outputs=[audio_input, transcript_output, soap_output, timing_output],
    )

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS, footer_links=[])