Spaces:
Running on Zero
Running on Zero
| import spaces # MUST be imported before torch | |
| import time | |
| import json | |
| import logging | |
| import uuid | |
| from datetime import datetime, timezone | |
| import torch | |
| import gradio as gr | |
| from pydantic import BaseModel | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from peft import PeftModel | |
| # ============================================================ | |
| # Logging setup (monitoring) | |
| # ============================================================ | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s | %(levelname)s | %(message)s", | |
| ) | |
| logger = logging.getLogger("mediguide") | |
| REQUEST_LOG_FILE = "request_log.jsonl" # ephemeral storage on free Spaces, resets on restart | |
| def log_request(request_id, message, response, latency_s, flagged, error=None): | |
| entry = { | |
| "request_id": request_id, | |
| "timestamp": datetime.now(timezone.utc).isoformat(), | |
| "message_len": len(message), | |
| "response_len": len(response) if response else 0, | |
| "latency_s": round(latency_s, 3), | |
| "flagged": flagged, | |
| "error": error, | |
| } | |
| logger.info(json.dumps(entry)) | |
| try: | |
| with open(REQUEST_LOG_FILE, "a") as f: | |
| f.write(json.dumps(entry) + "\n") | |
| except Exception as e: | |
| logger.warning(f"Could not write to request log file: {e}") | |
| # ============================================================ | |
| # Guardrails | |
| # ============================================================ | |
| EMERGENCY_KEYWORDS = [ | |
| "suicide", "kill myself", "want to die", "end my life", | |
| "chest pain", "can't breathe", "cannot breathe", "heart attack", | |
| "overdose", "severe bleeding", "unconscious", "stroke symptoms", | |
| ] | |
| DISCLAIMER = ( | |
| "\n\n⚠️ This is not a substitute for professional medical advice. " | |
| "If this is urgent, please contact a doctor or emergency services." | |
| ) | |
| EMERGENCY_RESPONSE = ( | |
| "This sounds like it could be a medical emergency. Please contact " | |
| "your local emergency services or go to the nearest emergency room " | |
| "immediately. I'm not able to provide emergency medical care." | |
| ) | |
| def check_emergency(message: str) -> bool: | |
| msg_lower = message.lower() | |
| return any(kw in msg_lower for kw in EMERGENCY_KEYWORDS) | |
| def apply_output_guardrails(response: str) -> str: | |
| if "not a substitute" not in response.lower() and "professional medical" not in response.lower(): | |
| response = response.strip() + DISCLAIMER | |
| return response | |
| # ============================================================ | |
| # Configuration | |
| # ============================================================ | |
| BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct" | |
| ADAPTER_MODEL = "rolmaxx/MediGuide-QLoRA" | |
| SYSTEM_PROMPT = """You are MediGuide, a medical conversational assistant. | |
| Provide clear, helpful, and appropriately cautious responses. | |
| Do not claim to replace a qualified healthcare professional. | |
| If symptoms may indicate an emergency, advise the user to seek | |
| appropriate professional medical care.""" | |
| # ============================================================ | |
| # Tokenizer & Model – loaded on CPU (no GPU yet) | |
| # ============================================================ | |
| logger.info("Loading tokenizer...") | |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| logger.info("Loading base model on CPU...") | |
| model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| torch_dtype=torch.float16, | |
| device_map="cpu", | |
| ) | |
| logger.info("Base model loaded on CPU.") | |
| logger.info("Loading MediGuide QLoRA adapter on CPU...") | |
| model = PeftModel.from_pretrained( | |
| model, | |
| ADAPTER_MODEL, | |
| torch_device="cpu", | |
| ) | |
| model = model.merge_and_unload() | |
| model.eval() | |
| logger.info("MediGuide loaded successfully (on CPU).") | |
| # ============================================================ | |
| # Core generation | |
| # ============================================================ | |
| def _run_model(messages): | |
| model.to("cuda") | |
| prompt = tokenizer.apply_chat_template( | |
| messages, | |
| tokenize=False, | |
| add_generation_prompt=True, | |
| ) | |
| inputs = tokenizer( | |
| prompt, | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=4096, | |
| ) | |
| inputs = {k: v.to("cuda") for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=256, | |
| do_sample=True, | |
| temperature=0.7, | |
| top_p=0.9, | |
| repetition_penalty=1.1, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:] | |
| response = tokenizer.decode(generated_tokens, skip_special_tokens=True) | |
| model.to("cpu") | |
| torch.cuda.empty_cache() | |
| return response.strip() | |
| def generate_response(message, history): | |
| """ | |
| history: list of {"role": ..., "content": ...} dicts (Gradio 'messages' format, | |
| also used as the request body shape for the /chat API below). | |
| """ | |
| request_id = str(uuid.uuid4())[:8] | |
| start = time.time() | |
| if not isinstance(message, str): | |
| message = str(message) if message is not None else "" | |
| message = message.strip() | |
| if not message: | |
| return "Please enter a valid question." | |
| if check_emergency(message): | |
| latency = time.time() - start | |
| log_request(request_id, message, EMERGENCY_RESPONSE, latency, flagged="emergency_keyword") | |
| return EMERGENCY_RESPONSE | |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}] | |
| for turn in history: | |
| if isinstance(turn, dict) and turn.get("role") in ("user", "assistant") and turn.get("content"): | |
| messages.append({"role": turn["role"], "content": str(turn["content"])}) | |
| messages.append({"role": "user", "content": message}) | |
| try: | |
| response = _run_model(messages) | |
| response = apply_output_guardrails(response) | |
| latency = time.time() - start | |
| log_request(request_id, message, response, latency, flagged=None) | |
| return response | |
| except Exception as e: | |
| latency = time.time() - start | |
| log_request(request_id, message, "", latency, flagged=None, error=str(e)) | |
| logger.error(f"[{request_id}] generation failed: {e}") | |
| return "Sorry, something went wrong while generating a response. Please try again." | |
| # ============================================================ | |
| # Gradio UI | |
| # ============================================================ | |
| with gr.Blocks(title="MediGuide") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🩺 MediGuide | |
| ### QLoRA Fine‑Tuned Medical Conversational Assistant | |
| MediGuide is based on **Qwen2.5‑1.5B‑Instruct** | |
| and fine‑tuned using **QLoRA** on medical dialogue data. | |
| ⚠️ **Disclaimer** | |
| This demo is for research and educational purposes only. | |
| It is **not a substitute** for professional medical advice, | |
| diagnosis, or treatment. | |
| """ | |
| ) | |
| chatbot = gr.Chatbot(label="MediGuide", height=500) | |
| message = gr.Textbox( | |
| label="Your question", | |
| placeholder="e.g. What are common symptoms of the flu?", | |
| lines=2, | |
| ) | |
| with gr.Row(): | |
| send_btn = gr.Button("Send", variant="primary") | |
| clear_btn = gr.Button("Clear") | |
| def respond(message, history): | |
| response = generate_response(message, history) | |
| history = history + [ | |
| {"role": "user", "content": message}, | |
| {"role": "assistant", "content": response}, | |
| ] | |
| return "", history | |
| message.submit(respond, [message, chatbot], [message, chatbot]) | |
| send_btn.click(respond, [message, chatbot], [message, chatbot]) | |
| clear_btn.click(lambda: [], outputs=chatbot) | |
| # ============================================================ | |
| # Launch via Gradio SDK (required for ZeroGPU detection), | |
| # then bolt a REST API onto Gradio's own FastAPI app. | |
| # ============================================================ | |
| class ChatRequest(BaseModel): | |
| message: str | |
| history: list = [] | |
| class ChatResponse(BaseModel): | |
| response: str | |
| demo.queue().launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| prevent_thread_lock=True, # let us keep running code after launch() returns | |
| ) | |
| def health(): | |
| return {"status": "ok"} | |
| def chat_endpoint(req: ChatRequest): | |
| response = generate_response(req.message, req.history) | |
| return ChatResponse(response=response) | |
| demo.block_thread() # keep the process alive (replaces the old blocking demo.launch()) |