app.py

"""
PlainScript — Patient-Friendly Medical Rewriter (Hugging Face ZeroGPU).

Designed specifically for ZeroGPU Spaces:
  - The GPU-using function is decorated with @spaces.GPU (a GPU is attached
    only while it runs, then released).
  - The model is instantiated at module scope and moved to CUDA eagerly;
    ZeroGPU maps the device transparently.
  - requirements.txt does NOT pin torch / gradio / spaces — the ZeroGPU base
    image provides those, and pinning them causes dependency conflicts.

Before ("base") and after ("fine-tuned") outputs come from a single
PeftModel by toggling the LoRA adapter with disable_adapter().

Portions of this file were developed with assistance from Claude (Anthropic).
"""

import os

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel


BASE_MODEL = "Qwen/Qwen2-1.5B-Instruct"
ADAPTER_ID = os.environ.get("ADAPTER_ID", "zkmine/plainscript-adapter")
HF_TOKEN = os.environ.get("HF_TOKEN")  # optional; only needed if adapter is private

SYSTEM_PROMPT = (
    "You are a medical text simplifier. Rewrite the following medical text "
    "into plain language that a patient with no medical background can "
    "understand. Preserve all key findings and conclusions. Do not add "
    "information not present in the original text."
)

EXAMPLES = [
    "A meta-analysis of randomized controlled trials demonstrated a "
    "statistically significant reduction in glycated hemoglobin (HbA1c) "
    "levels (mean difference -0.5%, 95% CI -0.7 to -0.3, p<0.001) in "
    "patients receiving the intervention compared to placebo.",
    "The systematic review found moderate-certainty evidence that cognitive "
    "behavioural therapy reduces the severity of chronic insomnia symptoms "
    "compared with treatment as usual, measured by the Pittsburgh Sleep "
    "Quality Index at 8 weeks post-intervention.",
    "Prophylactic administration of low-molecular-weight heparin was "
    "associated with reduced incidence of venous thromboembolism in "
    "post-operative orthopaedic patients (RR 0.50, 95% CI 0.33 to 0.76), "
    "though with a concomitant increase in minor bleeding events.",
]

# --- On ZeroGPU, torch is patched at import time and there is NO GPU at
#     module scope, so we cannot load model weights onto CUDA here. But we CAN
#     pre-download the files to the local cache at startup, so the GPU call
#     only has to load from disk (fast) and stays within the 60s GPU budget. ---
from huggingface_hub import snapshot_download

print("Loading tokenizer and pre-downloading model files...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

# Warm the HF cache so weights are on local disk before any GPU call.
snapshot_download(BASE_MODEL)
snapshot_download(ADAPTER_ID, token=HF_TOKEN)
print("Tokenizer ready and weights cached; model loads on first request.")

_model = None  # lazily populated inside the @spaces.GPU function


def _load_model():
    """
    Load base + LoRA adapter from the local cache and place on CUDA. Called
    once, lazily, from inside the @spaces.GPU function where a GPU is
    attached. Because files are already cached, this is just a disk load.
    """
    global _model
    if _model is not None:
        return _model

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        dtype=torch.float16,
        trust_remote_code=True,
    )
    peft_model = PeftModel.from_pretrained(base_model, ADAPTER_ID, token=HF_TOKEN)
    peft_model = peft_model.to("cuda")
    peft_model.eval()
    _model = peft_model
    return _model


def _build_prompt(text: str) -> str:
    """Render the chat prompt for a single user input."""
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": text},
    ]
    return tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )


def _decode(output_ids, input_len: int) -> str:
    """Decode only the newly generated tokens."""
    gen = output_ids[0][input_len:]
    return tokenizer.decode(gen, skip_special_tokens=True).strip()


@spaces.GPU(duration=60)
def rewrite(text: str):
    """
    Produce base ('before') and fine-tuned ('after') rewrites.

    Runs on a ZeroGPU-attached GPU. The base output is generated with the
    LoRA adapter disabled; the fine-tuned output with it enabled — one model
    in memory, toggled per call.

    Args:
        text: The medical text to rewrite.

    Returns:
        Tuple of (before_text, after_text).
    """
    if not text or not text.strip():
        return "", "Please paste some medical text above first."

    # Load the model on first use — we are inside @spaces.GPU here, so a GPU
    # is attached and CUDA is available (unlike at module scope).
    model = _load_model()

    prompt = _build_prompt(text)
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    input_len = inputs["input_ids"].shape[1]

    gen_kwargs = dict(
        max_new_tokens=400,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
    )

    with torch.no_grad():
        with model.disable_adapter():
            before_ids = model.generate(**inputs, **gen_kwargs)
        after_ids = model.generate(**inputs, **gen_kwargs)

    return _decode(before_ids, input_len), _decode(after_ids, input_len)


# --- UI ---

CUSTOM_CSS = """
.gradio-container {
  background: linear-gradient(135deg, #F0FDFA 0%, #E0F2FE 50%, #FCE7F3 100%);
  background-size: 400% 400%;
  animation: bgshift 18s ease infinite;
  max-width: 1080px !important; margin: auto !important;
}
@keyframes bgshift {
  0% { background-position: 0% 50%; }
  50% { background-position: 100% 50%; }
  100% { background-position: 0% 50%; }
}
#sticker-layer { position: fixed; inset: 0; pointer-events: none; z-index: 0; overflow: hidden; }
.sticker { position: absolute; font-size: 2.4rem; opacity: 0.13; animation: float 15s ease-in-out infinite; }
.sticker:nth-child(1){left:6%;top:18%;animation-delay:0s}
.sticker:nth-child(2){left:85%;top:12%;animation-delay:2s}
.sticker:nth-child(3){left:12%;top:73%;animation-delay:4s}
.sticker:nth-child(4){left:80%;top:68%;animation-delay:1s}
.sticker:nth-child(5){left:46%;top:8%;animation-delay:3s}
.sticker:nth-child(6){left:90%;top:44%;animation-delay:5s}
@keyframes float { 0%,100%{transform:translateY(0) rotate(-4deg)} 50%{transform:translateY(-24px) rotate(4deg)} }
.hero { position: relative; z-index: 1; text-align: center; padding: 24px 16px 4px; }
.hero h1 { font-size: 2.4rem; font-weight: 800; letter-spacing: -0.02em; color: #0F172A; margin: 0; }
.hero h1 .accent { color: #0B7C7B; }
.hero p { color: #475569; font-size: 1.02rem; margin: 8px auto 0; max-width: 600px; }
.ecg { width: 210px; height: 38px; margin: 12px auto 0; display: block; }
.ecg path { fill: none; stroke: #FB7185; stroke-width: 2.5; stroke-dasharray: 300; stroke-dashoffset: 300; animation: trace 2.2s linear infinite; }
@keyframes trace { to { stroke-dashoffset: -300; } }
#after_box textarea { background: rgba(240,253,250,0.9) !important; border-left: 4px solid #0EA5A4 !important; font-weight: 500; }
#before_box textarea { background: rgba(255,255,255,0.6) !important; border-left: 4px solid #94A3B8 !important; }
.disclaimer { text-align:center; color:#64748B; font-size:0.82rem; margin-top:16px; position:relative; z-index:1; }
"""

STICKERS = """
<div id="sticker-layer">
  <div class="sticker">🩺</div><div class="sticker">💊</div><div class="sticker">❤️</div>
  <div class="sticker">🩹</div><div class="sticker">💉</div><div class="sticker">🏥</div>
</div>
"""

HERO = """
<div class="hero">
  <h1>Plain<span class="accent">Script</span></h1>
  <p>Paste a dense medical abstract and watch it become language a patient can
     understand — the same model, before and after fine-tuning.</p>
  <svg class="ecg" viewBox="0 0 220 40">
    <path d="M0,20 L60,20 L70,20 L78,4 L88,36 L98,20 L120,20 L128,12 L136,28 L146,20 L220,20"/>
  </svg>
</div>
"""


def build_app() -> gr.Blocks:
    """Construct the Gradio Blocks UI."""
    with gr.Blocks(css=CUSTOM_CSS, title="PlainScript") as demo:
        gr.HTML(STICKERS)
        gr.HTML(HERO)

        inp = gr.Textbox(
            label="Medical text",
            placeholder="Paste a medical abstract or clinical summary here…",
            lines=6,
        )
        btn = gr.Button("Rewrite it", variant="primary")

        with gr.Row():
            before = gr.Textbox(label="Before · base Qwen2-1.5B", lines=9,
                                interactive=False, elem_id="before_box")
            after = gr.Textbox(label="After · fine-tuned on Cochrane", lines=9,
                               interactive=False, elem_id="after_box")

        gr.Examples(examples=EXAMPLES, inputs=inp)
        gr.HTML(
            '<div class="disclaimer">Drafts for review only — not medical '
            "advice. Outputs may contain errors and must be checked by a "
            "qualified professional.</div>"
        )

        btn.click(fn=rewrite, inputs=inp, outputs=[before, after])

    return demo


if __name__ == "__main__":
    build_app().launch()