"""
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 = """