src/simplify.py

"""
simplify.py — turn clinical text into plain, 6th-grade language
===============================================================
Two modes:
  * LLM mode (production): an instruction-tuned model rewrites for fluency at a
    target grade level, told explicitly to PRESERVE every medical fact (drugs,
    doses, dates, warnings) and to ADD nothing.
  * Deterministic mode (fallback + safety layer): expand dangerous abbreviations,
    swap jargon for plain words, and shorten sentences. Always run, even with the
    LLM, so abbreviations like "BID" or "PO" are never left ambiguous in audio.
"""
from __future__ import annotations
import re
from pathlib import Path
import yaml

GLOSSARY_PATH = Path(__file__).resolve().parent.parent / "config" / "medical_glossary.yaml"


def load_glossary(path: Path = GLOSSARY_PATH) -> dict:
    with open(path) as f:
        return yaml.safe_load(f)


def expand_abbreviations(text: str, gloss: dict) -> str:
    # CASE-SENSITIVE: "Pt" (you) must not collide with "PT" (physical therapy).
    # Longest keys first so multi-letter forms win.
    for abbr in sorted(gloss.get("abbreviations", {}), key=len, reverse=True):
        full = gloss["abbreviations"][abbr]
        text = re.sub(rf"(?<![A-Za-z]){re.escape(abbr)}(?![A-Za-z])", full, text)
    return text


def replace_jargon(text: str, gloss: dict) -> str:
    # longest phrases first so "high blood pressure" wins over "pressure"
    for term in sorted(gloss.get("jargon", {}), key=len, reverse=True):
        plain = gloss["jargon"][term]
        text = re.sub(re.escape(term), plain, text, flags=re.IGNORECASE)
    return text


def shorten_sentences(text: str, max_words: int = 14) -> str:
    """Split long sentences at conjunctions/semicolons into shorter ones."""
    out = []
    for sent in re.split(r"(?<=[.!?])\s+", text.strip()):
        words = sent.split()
        if len(words) <= max_words:
            out.append(sent)
            continue
        parts = re.split(r"\s*(?:;|,\s+(?:and|but|because|so)\b)\s*", sent)
        for p in parts:
            p = p.strip().rstrip(".")
            if p:
                out.append(p[0].upper() + p[1:] + ".")
    return " ".join(out)


def _capitalize_sentences(text: str) -> str:
    # Capitalize the first letter of each sentence for readable output.
    def up(m): return m.group(1) + m.group(2).upper()
    text = re.sub(r"(^|[.!?]\s+)([a-z])", up, text)
    return text


def deterministic_simplify(text: str, gloss: dict | None = None) -> str:
    gloss = gloss or load_glossary()
    text = expand_abbreviations(text, gloss)
    text = replace_jargon(text, gloss)
    text = shorten_sentences(text)
    text = re.sub(r"\s+", " ", text).strip()
    return _capitalize_sentences(text)


_LLM_CACHE: dict = {}


def _get_llm_pipeline(model_id: str):
    """Load-once cache: without this, every simplify() call reloaded the whole
    model from disk (~10-15s of pure loading overhead on top of generation
    time) -- fine for a one-shot script, but not for an interactive Space
    where the same process handles many requests."""
    if model_id not in _LLM_CACHE:
        from transformers import pipeline
        _LLM_CACHE[model_id] = pipeline("text-generation", model=model_id, dtype="auto")
    return _LLM_CACHE[model_id]


LLM_SYSTEM_PROMPT = """You rewrite medical instructions in plain language for patients. Rules you MUST follow:
- Copy every NUMBER exactly, character for character: doses (e.g. "25 mg"), how many times a day, and dates. Never rewrite a dose as just "your medicine" or "as directed" -- always repeat the exact number and unit.
- Keep every medicine name, dose, how often, date, and warning sign from the original.
- Do NOT add any new advice, diagnosis, or information that is not in the text.
- Do NOT add commentary, notes, disclaimers, or anything about your own process. Output ONLY the rewritten text.
- Use short sentences and everyday words for a {grade}th-grade reader.
- Speak DIRECTLY TO THE PATIENT using "you" and "your" (e.g. "You had a heart attack" and "Take your medicine", never "the patient" or "they").
- Expand abbreviations (for example, "BID" becomes "two times a day")."""


def llm_simplify(text: str, model_id: str, grade: int = 6, max_new_tokens: int = 300) -> str:
    """Production path. Requires `transformers` (+ a model). Kept optional so the
    package imports and runs without heavy deps.

    Uses the model's chat template (via the `messages=` pipeline argument) rather
    than a raw text-continuation prompt: an early version of this function used a
    plain string prompt with a hand-rolled "Plain-language version:" marker, which
    on a chat-tuned model (Qwen2.5-Instruct) produced hundreds of tokens of
    unstoppable rambling (meta-commentary, a second re-explanation "for a
    6th-grader", emoji, sign-offs) because the model never saw a proper end-of-
    turn signal. The chat template gives the model a real stopping point.
    """
    gen = _get_llm_pipeline(model_id)
    messages = [
        {"role": "system", "content": LLM_SYSTEM_PROMPT.format(grade=grade)},
        {"role": "user", "content": text},
    ]
    out = gen(messages, max_new_tokens=max_new_tokens, do_sample=False)[0]["generated_text"]
    return out[-1]["content"].strip()


def simplify(text: str, use_llm: bool = False, model_id: str | None = None,
             grade: int = 6, gloss: dict | None = None) -> str:
    gloss = gloss or load_glossary()
    if use_llm and model_id:
        try:
            simplified = llm_simplify(text, model_id, grade)
            # always run the deterministic safety layer on top (expand abbrevs)
            return deterministic_simplify(simplified, gloss)
        except Exception as e:  # fall back gracefully
            print(f"[simplify] LLM unavailable ({e}); using deterministic fallback.")
    return deterministic_simplify(text, gloss)