Spaces:
Sleeping
Sleeping
| """ZeroGPU Gradio Space — crisis triage student model demo, with voice in/out | |
| and a model switcher between trained adapters. | |
| ZeroGPU only attaches a real GPU for the duration of a function decorated | |
| with @spaces.GPU -- at module/import time there is no GPU visible, so every | |
| model must load on CPU first and move to "cuda" only inside that function. | |
| Models are LAZY-LOADED: nothing but the dropdown's default entry loads at | |
| startup. Switching the dropdown loads that model on CPU (cached after first | |
| use) and frees the previous one, so you never hold a 1.5B and a 7B model in | |
| CPU RAM at the same time. STT/Whisper and TTS/VITS are always-on and shared | |
| across model choices. | |
| Pipeline per click (all inside ONE @spaces.GPU call, to use a single GPU | |
| borrow instead of three): | |
| voice in -> Whisper STT -> transcript | |
| transcript -> triage adapter (selected model) -> structured JSON | |
| JSON["plain_language_summary"] -> TTS -> spoken response | |
| EDIT this dict to add/remove selectable models: | |
| MODEL_REGISTRY | |
| """ | |
| import spaces # must be imported before torch/transformers/peft on ZeroGPU | |
| import gc | |
| import html | |
| import json | |
| import re | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from peft import PeftModel | |
| from transformers import ( | |
| AutoModelForCausalLM, | |
| AutoTokenizer, | |
| VitsModel, | |
| WhisperForConditionalGeneration, | |
| WhisperProcessor, | |
| ) | |
| MODEL_REGISTRY = { | |
| "Qwen2.5-1.5B (fast)": { | |
| "base": "Qwen/Qwen2.5-1.5B-Instruct", | |
| "adapter": "NathanPereira/crisis-triage-adapter", | |
| }, | |
| "Qwen3-4B (balanced)": { | |
| "base": "Qwen/Qwen3-4B", | |
| "adapter": "NathanPereira/crisis-triage-adapter-qwen3-4b", | |
| }, | |
| "DeepSeek-R1-Distill-Qwen-7B (stronger)": { | |
| "base": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", | |
| "adapter": "NathanPereira/crisis-triage-adapter-r1-distill-7b", | |
| }, | |
| } | |
| DEFAULT_MODEL_NAME = "Qwen2.5-1.5B (fast)" | |
| ASR_MODEL_ID = "openai/whisper-small" | |
| TTS_MODEL_ID = "facebook/mms-tts-eng" | |
| SYSTEM_PROMPT = ( | |
| "You are a research-only crisis triage assistant. " | |
| "Return only valid JSON matching the triage schema. " | |
| "Use exact evidence spans copied from the input text. " | |
| "Do not diagnose, prescribe medication, or provide therapy instructions." | |
| ) | |
| # name -> (tokenizer, model) on CPU; populated lazily, at most one entry at a time | |
| _loaded = {} | |
| def _get_triage_model(name: str): | |
| if name in _loaded: | |
| return _loaded[name] | |
| # evict whatever was loaded before so we never hold two student models in RAM | |
| if _loaded: | |
| old_name = next(iter(_loaded)) | |
| del _loaded[old_name] | |
| gc.collect() | |
| cfg = MODEL_REGISTRY[name] | |
| print(f"loading triage model '{name}' (CPU) ...") | |
| tok = AutoTokenizer.from_pretrained(cfg["base"]) | |
| base = AutoModelForCausalLM.from_pretrained(cfg["base"], torch_dtype=torch.float16) | |
| model = PeftModel.from_pretrained(base, cfg["adapter"], torch_device="cpu").eval() | |
| _loaded[name] = (tok, model) | |
| return tok, model | |
| print("loading Whisper STT (CPU) ...") | |
| asr_processor = WhisperProcessor.from_pretrained(ASR_MODEL_ID) | |
| asr_model = WhisperForConditionalGeneration.from_pretrained(ASR_MODEL_ID, torch_dtype=torch.float16).eval() | |
| print("loading TTS (CPU) ...") | |
| tts_tok = AutoTokenizer.from_pretrained(TTS_MODEL_ID) | |
| tts_model = VitsModel.from_pretrained(TTS_MODEL_ID).eval() | |
| print(f"warming default triage model: {DEFAULT_MODEL_NAME} ...") | |
| _get_triage_model(DEFAULT_MODEL_NAME) | |
| print("all models ready") | |
| def _salvage_tier(raw: str) -> int | None: | |
| m = re.search(r'"risk_tier"\s*:\s*([0-3])', raw) | |
| return int(m.group(1)) if m else None | |
| def _extract_json(raw: str) -> dict: | |
| if "</think>" in raw: | |
| raw = raw.split("</think>", 1)[1] | |
| s, e = raw.find("{"), raw.rfind("}") | |
| if s < 0 or e <= s: | |
| raise ValueError("no JSON object") | |
| return json.loads(raw[s : e + 1]) | |
| def _transcribe(audio) -> str: | |
| # gradio Audio(type="numpy") gives (sample_rate, np.ndarray); no file I/O, no librosa. | |
| import torchaudio | |
| sr, data = audio | |
| wav = torch.as_tensor(data).float() | |
| if wav.ndim == 2: # stereo -> mono | |
| wav = wav.mean(dim=1) | |
| if data.dtype.kind in "iu": # int PCM -> [-1, 1] | |
| wav = wav / float(np.iinfo(data.dtype).max) | |
| if sr != 16000: | |
| wav = torchaudio.functional.resample(wav, sr, 16000) | |
| inputs = asr_processor(wav.numpy(), sampling_rate=16000, return_tensors="pt") | |
| input_features = inputs.input_features.to("cuda", dtype=torch.float16) | |
| with torch.no_grad(): | |
| pred_ids = asr_model.generate(input_features, language="en", task="transcribe") | |
| return asr_processor.batch_decode(pred_ids, skip_special_tokens=True)[0].strip() | |
| def _normalize_tier(tier): | |
| """Collapse the many risk_tier shapes the model emits into an int 0..3. | |
| Handles ints, floats, numeric strings ("2"), and free-text labels | |
| ("High Risk", "imminent", "low") so downstream phrasing is consistent. | |
| Returns None if nothing recognizable. | |
| """ | |
| if isinstance(tier, bool) or tier is None: | |
| return None | |
| if isinstance(tier, (int, float)): | |
| t = int(tier) | |
| return t if 0 <= t <= 3 else None | |
| if isinstance(tier, str): | |
| s = tier.strip().lower() | |
| m = re.search(r"[0-3]", s) | |
| if m: | |
| return int(m.group()) | |
| if any(w in s for w in ("imminent", "critical", "emergency", "severe", "tier 3")): | |
| return 3 | |
| if "high" in s: | |
| return 2 | |
| if any(w in s for w in ("moderate", "medium", "elevated")): | |
| return 1 | |
| if any(w in s for w in ("low", "minimal", "none", "no risk")): | |
| return 0 | |
| return None | |
| def _first_text(value) -> str: | |
| """Pull a usable sentence out of a str / list / dict field.""" | |
| if isinstance(value, str): | |
| return value.strip() | |
| if isinstance(value, (list, tuple)): | |
| for v in value: | |
| t = _first_text(v) | |
| if t: | |
| return t | |
| if isinstance(value, dict): | |
| for v in value.values(): | |
| t = _first_text(v) | |
| if t: | |
| return t | |
| return "" | |
| # tier-aware openers, written to be spoken aloud like a person talking, not a report. | |
| # each has a connector for folding a next-step in naturally instead of a labeled clause. | |
| _TIER_VOICE = { | |
| 0: dict( | |
| opener="Thanks for sharing that with me. From what you've described, things " | |
| "don't sound like they're at a crisis point right now.", | |
| connector="That said, it can still help to", | |
| ), | |
| 1: dict( | |
| opener="I can hear that you're going through a tough time. What you're " | |
| "describing sounds like real distress, and it's worth taking seriously.", | |
| connector="One thing that might help is to", | |
| ), | |
| 2: dict( | |
| opener="I'm genuinely concerned about what you've shared. This sounds serious, " | |
| "and I want you to know your safety matters right now.", | |
| connector="Please", | |
| ), | |
| 3: dict( | |
| opener="What you've told me sounds like an emergency, and I'm worried about " | |
| "your safety right now.", | |
| connector="Right now, please", | |
| ), | |
| } | |
| _DEFAULT_VOICE = dict( | |
| opener="Thank you for trusting me with this. Whatever you're going through, " | |
| "you don't have to face it alone.", | |
| connector="It could help to", | |
| ) | |
| # always-on safety line for higher tiers, regardless of what JSON fields exist | |
| _SAFETY_LINE = ("And if you're ever in immediate danger, please call your local " | |
| "emergency number, or reach out to a crisis line right away.") | |
| def _lowercase_lead(s: str) -> str: | |
| """Lowercase a sentence's first letter so it reads naturally mid-sentence.""" | |
| return s[:1].lower() + s[1:] if s else s | |
| def _strip_trailing_period(s: str) -> str: | |
| return s[:-1] if s.endswith(".") else s | |
| def _build_summary(parsed: dict) -> str: | |
| """Compose a warm, conversational, tier-aware spoken triage summary. | |
| Written for a synthetic chain-of-thought triage adapter: it talks to the person | |
| the way a caring listener would, folds in one concrete next step as part of the | |
| same flowing sentence rather than a labeled clause, and always closes high-risk | |
| tiers with a safety line -- even when the model drifts off-schema. | |
| Prefers the trained `plain_language_summary` verbatim when present. | |
| """ | |
| if not isinstance(parsed, dict): | |
| return "" | |
| direct = parsed.get("plain_language_summary") | |
| if isinstance(direct, str) and direct.strip(): | |
| return direct.strip() | |
| tier = _normalize_tier(parsed.get("risk_tier")) | |
| voice = _TIER_VOICE.get(tier, _DEFAULT_VOICE) | |
| sentence = voice["opener"] | |
| # fold one concrete next step into the same breath as the opener, instead of | |
| # a separate labeled sentence | |
| for key in ("recommended_next_step", "action_steps", "resources", | |
| "clinical_rationale", "risk_factors"): | |
| action = _first_text(parsed.get(key)) | |
| if action: | |
| action = _strip_trailing_period(_lowercase_lead(action.strip())) | |
| sentence = f"{sentence} {voice['connector']} {action}." | |
| break | |
| escalate = (parsed.get("escalation_required") in (True, "true", "True", "yes", 1) | |
| or (tier is not None and tier >= 2)) | |
| if escalate: | |
| sentence = f"{sentence} {_SAFETY_LINE}" | |
| return sentence.strip() | |
| # badge color/label per normalized tier, used to render the report card | |
| _TIER_META = { | |
| 0: dict(color="#2e7d32", bg="#e8f5e9", label="Tier 0 · Minimal"), | |
| 1: dict(color="#a16207", bg="#fef9c3", label="Tier 1 · Moderate"), | |
| 2: dict(color="#c2410c", bg="#ffedd5", label="Tier 2 · High"), | |
| 3: dict(color="#b91c1c", bg="#fee2e2", label="Tier 3 · Imminent"), | |
| } | |
| _TIER_META_UNKNOWN = dict(color="#374151", bg="#f3f4f6", label="Tier unclear") | |
| # keys rendered up front, in this order, with friendly labels + how to render them | |
| _REPORT_FIELDS = [ | |
| ("evidence_spans", "Evidence from the text", "quotes"), | |
| ("risk_factors", "Risk factors", "list"), | |
| ("protective_factors", "Protective factors", "list"), | |
| ("cssrs_axes", "C-SSRS axes", "table"), | |
| ("clinical_rationale", "Clinical rationale", "text"), | |
| ("recommended_next_step", "Recommended next step", "highlight"), | |
| ("uncertainty_flags", "Uncertainty flags", "list"), | |
| ] | |
| # fields shown elsewhere (tier badge, escalation banner) or internal -- skip in the body | |
| _REPORT_SKIP_KEYS = { | |
| "risk_tier", "plain_language_summary", "escalation_required", | |
| "_parse_error", "_raw", "confidence", | |
| } | |
| def _esc(value) -> str: | |
| return html.escape(str(value), quote=True) | |
| def _render_inline(value) -> str: | |
| """Render a str/list/dict value as a flowing HTML fragment (no own heading).""" | |
| if isinstance(value, str): | |
| return _esc(value) | |
| if isinstance(value, (list, tuple)): | |
| items = [_render_inline(v) for v in value if _render_inline(v)] | |
| if not items: | |
| return "" | |
| return "<ul class='ct-list'>" + "".join(f"<li>{i}</li>" for i in items) + "</ul>" | |
| if isinstance(value, dict): | |
| rows = "".join( | |
| f"<tr><td class='ct-k'>{_esc(k)}</td><td>{_render_inline(v)}</td></tr>" | |
| for k, v in value.items() if _render_inline(v) | |
| ) | |
| return f"<table class='ct-table'>{rows}</table>" if rows else "" | |
| if value in (None, "", [], {}): | |
| return "" | |
| return _esc(value) | |
| def _render_section(label: str, value, kind: str) -> str: | |
| if value in (None, "", [], {}): | |
| return "" | |
| if kind == "quotes" and isinstance(value, (list, tuple)): | |
| body = "".join(f"<blockquote class='ct-quote'>“{_esc(v)}”</blockquote>" for v in value if str(v).strip()) | |
| elif kind == "highlight": | |
| body = f"<div class='ct-highlight'>{_render_inline(value)}</div>" | |
| else: | |
| body = _render_inline(value) | |
| if not body: | |
| return "" | |
| return f"<div class='ct-section'><div class='ct-section-title'>{_esc(label)}</div>{body}</div>" | |
| def render_report(parsed: dict) -> str: | |
| """Render the triage JSON as a readable HTML report card instead of raw JSON. | |
| Walks the known schema fields first (evidence, rationale, next step, ...) in a | |
| fixed friendly order, then falls back to rendering any extra/unexpected keys the | |
| model produced (schema drift), so nothing the model says is silently dropped. | |
| """ | |
| if not isinstance(parsed, dict): | |
| return "<div class='ct-card'><i>No structured output to show.</i></div>" | |
| if parsed.get("_parse_error"): | |
| return ( | |
| "<div class='ct-card'><div class='ct-section'>" | |
| "<div class='ct-section-title'>Couldn't parse a structured response</div>" | |
| "<p>The model's raw output didn't contain valid JSON -- see the raw output box below.</p>" | |
| "</div></div>" | |
| ) | |
| tier = _normalize_tier(parsed.get("risk_tier")) | |
| meta = _TIER_META.get(tier, _TIER_META_UNKNOWN) | |
| confidence = parsed.get("confidence") | |
| conf_html = "" | |
| if isinstance(confidence, (int, float)): | |
| conf_html = f"<span class='ct-confidence'>confidence {confidence:.0%}" \ | |
| f"</span>" if confidence <= 1 else f"<span class='ct-confidence'>confidence {confidence}</span>" | |
| escalate = (parsed.get("escalation_required") in (True, "true", "True", "yes", 1) | |
| or (tier is not None and tier >= 2)) | |
| sections = "".join( | |
| _render_section(label, parsed.get(key), kind) for key, label, kind in _REPORT_FIELDS | |
| ) | |
| extra_keys = [k for k in parsed if k not in _REPORT_SKIP_KEYS | |
| and k not in {f[0] for f in _REPORT_FIELDS}] | |
| extras = "".join( | |
| _render_section(k.replace("_", " ").title(), parsed.get(k), "text") for k in extra_keys | |
| ) | |
| banner = ( | |
| "<div class='ct-escalate'>⚠ This case may need urgent attention or escalation.</div>" | |
| if escalate else "" | |
| ) | |
| return f""" | |
| <div class="ct-card"> | |
| <style> | |
| .ct-card {{ font-family: inherit; }} | |
| .ct-badge {{ display:inline-block; padding:4px 12px; border-radius:999px; | |
| font-weight:600; font-size:0.85em; color:{meta['color']}; | |
| background:{meta['bg']}; margin-right:8px; }} | |
| .ct-confidence {{ font-size:0.85em; opacity:0.75; }} | |
| .ct-escalate {{ margin-top:10px; padding:8px 12px; border-radius:8px; | |
| background:#fee2e2; color:#b91c1c; font-weight:600; font-size:0.9em; }} | |
| .ct-section {{ margin-top:16px; }} | |
| .ct-section-title {{ font-weight:600; font-size:0.85em; text-transform:uppercase; | |
| letter-spacing:0.03em; opacity:0.6; margin-bottom:6px; }} | |
| .ct-list {{ margin:0; padding-left:1.2em; }} | |
| .ct-quote {{ margin:4px 0; padding:6px 10px; border-left:3px solid #d1d5db; | |
| font-style:italic; opacity:0.85; }} | |
| .ct-highlight {{ padding:10px 12px; border-radius:8px; background:#eff6ff; | |
| border:1px solid #bfdbfe; }} | |
| .ct-table {{ width:100%; border-collapse:collapse; }} | |
| .ct-table td {{ padding:4px 8px; vertical-align:top; border-bottom:1px solid #e5e7eb; }} | |
| .ct-table td.ct-k {{ font-weight:600; white-space:nowrap; opacity:0.7; }} | |
| </style> | |
| <div><span class="ct-badge">{_esc(meta['label'])}</span>{conf_html}</div> | |
| {banner} | |
| {sections} | |
| {extras} | |
| </div> | |
| """ | |
| def _synthesize(text: str) -> tuple[int, np.ndarray]: | |
| inputs = tts_tok(text, return_tensors="pt").to("cuda") | |
| with torch.no_grad(): | |
| waveform = tts_model(**inputs).waveform | |
| audio = waveform.squeeze().to(torch.float32).cpu().numpy() | |
| return tts_model.config.sampling_rate, audio | |
| # borrows a real GPU only for the duration of this call | |
| def analyze(model_name: str, text: str, audio) -> tuple[str, str, str, str, str, tuple]: | |
| tok, model = _get_triage_model(model_name) | |
| asr_model.to("cuda") | |
| tts_model.to("cuda") | |
| model.to("cuda") | |
| transcript = "" | |
| if audio is not None: | |
| try: | |
| transcript = _transcribe(audio) | |
| text = transcript or text | |
| except Exception as exc: | |
| transcript = f"(speech-to-text failed: {exc})" | |
| if not text or not text.strip(): | |
| placeholder = "<div class='ct-card'><i>Enter text or record audio first.</i></div>" | |
| return placeholder, "", "", transcript, "", None | |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": text}] | |
| prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| inputs = tok(prompt, return_tensors="pt").to("cuda") | |
| with torch.no_grad(): | |
| out = model.generate( | |
| **inputs, max_new_tokens=400, do_sample=False, pad_token_id=tok.eos_token_id | |
| ) | |
| raw = tok.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True) | |
| try: | |
| parsed = _extract_json(raw) | |
| except Exception: | |
| parsed = {"risk_tier": _salvage_tier(raw), "_parse_error": True, "_raw": raw} | |
| summary = _build_summary(parsed) or ( | |
| "I couldn't generate a clear summary -- please check the structured output " | |
| "or contact a crisis line directly." | |
| ) | |
| try: | |
| speech_audio = _synthesize(summary) | |
| except Exception as exc: | |
| summary = f"{summary}\n\n(text-to-speech failed: {exc})" | |
| speech_audio = None | |
| return ( | |
| render_report(parsed), | |
| json.dumps(parsed, indent=2, ensure_ascii=False), | |
| raw, | |
| transcript, | |
| summary, | |
| speech_audio, | |
| ) | |
| with gr.Blocks(title="Crisis Triage — Research Demo") as demo: | |
| gr.Markdown( | |
| "## Crisis Triage — Research Demo\n" | |
| "Running on **ZeroGPU**, with voice input (Whisper STT), voice output (TTS), " | |
| "and a switchable student model.\n\n" | |
| "⚠ Research prototype. Not a clinical tool. " | |
| "Outputs are not a substitute for professional judgment or emergency services. " | |
| "If you or someone else is in immediate danger, contact local emergency services " | |
| "or a crisis line directly." | |
| ) | |
| model_choice = gr.Dropdown( | |
| label="Model", | |
| choices=list(MODEL_REGISTRY.keys()), | |
| value=DEFAULT_MODEL_NAME, | |
| info="Switching loads the new adapter on first use (a few seconds) and evicts the previous one.", | |
| ) | |
| with gr.Row(): | |
| inp = gr.Textbox(label="Text to assess", lines=5, placeholder="Type a message or transcript…") | |
| mic = gr.Audio(label="...or record/upload voice", sources=["microphone", "upload"], type="numpy") | |
| btn = gr.Button("Analyze", variant="primary") | |
| transcript_box = gr.Textbox(label="Transcribed speech (if audio was used)", lines=2, interactive=False) | |
| out_report = gr.HTML(label="Triage report") | |
| with gr.Accordion("Raw model output (JSON / debug)", open=False): | |
| out_json = gr.Code(label="Structured JSON output", language="json") | |
| out_raw = gr.Textbox(label="Raw model output", lines=4) | |
| out_summary = gr.Textbox(label="Plain-language summary (spoken aloud)", lines=2, interactive=False) | |
| out_audio = gr.Audio(label="Spoken response", type="numpy") | |
| btn.click( | |
| analyze, | |
| inputs=[model_choice, inp, mic], | |
| outputs=[out_report, out_json, out_raw, transcript_box, out_summary, out_audio], | |
| ) | |
| demo.launch() |