llm.py

"""Audio -> MedASR transcript -> MedGemma 4B SOAP note pipeline."""

import os
import re
import time
from dataclasses import dataclass

import spaces
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor, pipeline

ASR_MODEL_ID = os.environ.get("MEDASR_MODEL_ID", "google/medasr")
LLM_MODEL_ID = os.environ.get("MEDGEMMA_MODEL_ID", "google/medgemma-4b-it")

SYSTEM_PROMPT = (
    "You are a clinical documentation engine. Convert the transcript into a "
    "SOAP note (Subjective, Objective, Assessment, Plan). Infer which "
    "statements come from the doctor versus the patient based on context "
    "(questions, clinical observations vs. symptom descriptions) and "
    "attribute them accordingly (e.g. 'Patient reports...', 'On "
    "questioning by the physician...'). Do not fabricate any detail not "
    "present in the transcript. If speaker attribution is unclear, mark it "
    "as unclear rather than guessing confidently."
)

# 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. 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
_llm_processor = None


def _device_and_dtype():
    if torch.cuda.is_available():
        return "cuda", torch.bfloat16
    if torch.backends.mps.is_available():
        return "mps", torch.float16
    return "cpu", torch.float32


def get_asr_pipeline():
    global _asr_pipe
    if _asr_pipe is None:
        device, dtype = _device_and_dtype()
        _asr_pipe = pipeline(
            "automatic-speech-recognition",
            model=ASR_MODEL_ID,
            device=device,
            dtype=dtype,
        )
    return _asr_pipe


def get_llm():
    global _llm_model, _llm_processor
    if _llm_model is None:
        device, dtype = _device_and_dtype()
        _llm_model = AutoModelForImageTextToText.from_pretrained(
            LLM_MODEL_ID, dtype=dtype, device_map=device
        )
        _llm_processor = AutoProcessor.from_pretrained(LLM_MODEL_ID)
    return _llm_model, _llm_processor


def transcribe(audio_path: str) -> str:
    """Transcribe an audio file to text using MedASR, chunked for long audio."""
    pipe = get_asr_pipeline()
    result = pipe(audio_path, chunk_length_s=20, stride_length_s=2)
    return result["text"].strip()


def generate_soap_note(transcript: str) -> str:
    model, processor = get_llm()
    messages = [
        {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
        {"role": "user", "content": [{"type": "text", "text": transcript}]},
        {"role": "assistant", "content": [{"type": "text", "text": SOAP_PREFIX}]},
    ]
    inputs = processor.apply_chat_template(
        messages,
        continue_final_message=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    input_len = inputs["input_ids"].shape[-1]

    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=768,
            do_sample=True,
            temperature=0.2,
            top_p=0.9,
            repetition_penalty=1.15,
            no_repeat_ngram_size=3,
        )

    new_tokens = generated_ids[0][input_len:]
    completion = processor.decode(new_tokens, skip_special_tokens=True)
    return _normalize_soap_note(SOAP_PREFIX + completion)


@dataclass
class PipelineResult:
    transcript: str
    soap_note: str
    transcription_seconds: float
    generation_seconds: float


@spaces.GPU(duration=120)
def run_pipeline(audio_path: str) -> PipelineResult:
    t0 = time.perf_counter()
    transcript = transcribe(audio_path)
    t1 = time.perf_counter()

    soap_note = generate_soap_note(transcript)
    t2 = time.perf_counter()

    return PipelineResult(
        transcript=transcript,
        soap_note=soap_note,
        transcription_seconds=t1 - t0,
        generation_seconds=t2 - t1,
    )