Create app.py

app.py

@@ -0,0 +1,159 @@
+# app/main.py
+
+import os
+import time
+import logging
+from typing import Optional
+
+from fastapi import FastAPI, HTTPException, Query
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("biogpt_chatbot")
+
+# =========================
+# PROMPT TEMPLATES
+# =========================
+MEDICAL_PROMPTS = {
+    "dermatology": """
+You are DermX-AI, a specialized medical AI assistant trained in dermatology. 
+Your role is to provide clear, evidence-based information about skin conditions, 
+diagnostic insights, and treatment options. 
+
+- Use simple but professional language, suitable for both patients and clinicians.  
+- When explaining, balance medical accuracy with user-friendly clarity.  
+- For any uncertain or critical cases, clearly advise consultation with a dermatologist.  
+- Always include safety reminders and disclaimers.  
+""",
+    "general": """
+You are a medical AI assistant designed to provide helpful, evidence-based health information.  
+When answering:  
+- Ensure accuracy and clarity in medical explanations.  
+- Provide actionable lifestyle and preventive care suggestions where applicable.  
+- Avoid giving definitive diagnoses or prescriptions—always emphasize professional medical consultation.  
+- Be empathetic, supportive, and professional in tone.  
+""",
+    "disclaimer": """
+⚠️ Important: I am an AI medical assistant, not a licensed healthcare professional. 
+The information provided is for educational purposes only and should not be 
+considered a substitute for professional medical advice, diagnosis, or treatment.  
+Please consult a dermatologist or qualified healthcare provider for personalized care.  
+""",
+}
+
+# =========================
+# FASTAPI SETUP
+# =========================
+class ChatRequest(BaseModel):
+    question: str
+    context: Optional[str] = None
+    mode: Optional[str] = "dermatology"  # "dermatology" | "general"
+    max_new_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+
+class ChatResponse(BaseModel):
+    answer: str
+    model: str
+    took_seconds: float
+    confidence: int
+    sources: list
+
+app = FastAPI(title="BioGPT-Large Medical Chatbot")
+
+MODEL_ID = os.environ.get("MODEL_ID", "microsoft/BioGPT-Large")
+MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "200"))
+TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.7"))
+TOP_P = float(os.environ.get("TOP_P", "0.9"))
+DEVICE = int(os.environ.get("DEVICE", "-1"))  # -1 = CPU
+USE_4BIT = os.environ.get("USE_4BIT", "false").lower() == "true"
+
+generator = None
+
+@app.on_event("startup")
+def load_model():
+    global generator
+    try:
+        logger.info(f"Loading model: {MODEL_ID}")
+        if USE_4BIT:
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype="float16",
+                bnb_4bit_use_double_quant=True,
+            )
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
+                quantization_config=bnb_config,
+                device_map="auto",
+                trust_remote_code=True,
+            )
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+        generator = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device=DEVICE,
+        )
+        logger.info("Model loaded successfully.")
+    except Exception as e:
+        logger.exception("Model loading failed")
+        generator = None
+
+@app.post("/chat", response_model=ChatResponse)
+def chat(req: ChatRequest):
+    if generator is None:
+        raise HTTPException(status_code=500, detail="Model not available.")
+
+    if not req.question.strip():
+        raise HTTPException(status_code=400, detail="Question cannot be empty")
+
+    # Select system prompt
+    mode = req.mode.lower() if req.mode else "dermatology"
+    system_prompt = MEDICAL_PROMPTS.get(mode, MEDICAL_PROMPTS["general"])
+
+    # Build final prompt
+    prompt = f"{system_prompt}\n\nUser Question: {req.question.strip()}\n\nAI Answer:"
+    if req.context:
+        prompt = req.context.strip() + "\n\n" + prompt
+
+    max_new = req.max_new_tokens or MAX_NEW_TOKENS
+    temp = req.temperature or TEMPERATURE
+    top_p = req.top_p or TOP_P
+
+    logger.info(f"Generating answer for: {req.question[:80]}...")
+    t0 = time.time()
+
+    try:
+        outputs = generator(
+            prompt,
+            max_new_tokens=max_new,
+            temperature=temp,
+            top_p=top_p,
+            do_sample=True,
+            return_full_text=False,
+            num_return_sequences=1,
+        )
+        answer = outputs[0]["generated_text"].strip()
+
+        # Always append disclaimer
+        final_answer = f"{answer}\n\n{MEDICAL_PROMPTS['disclaimer']}"
+
+        took = time.time() - t0
+        confidence = min(95, 70 + int(len(answer) / 50))
+
+        return ChatResponse(
+            answer=final_answer,
+            model=MODEL_ID,
+            took_seconds=round(took, 2),
+            confidence=confidence,
+            sources=["HuggingFace", MODEL_ID],
+        )
+    except Exception as e:
+        logger.exception("Generation failed")
+        raise HTTPException(status_code=500, detail=str(e))