Update app.py

app.py

@@ -2,16 +2,30 @@ import os
 import time
 import logging
 from typing import Optional
+
+# =============================
+# Hugging Face cache fix for Spaces
+# =============================
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache/huggingface/transformers"
+os.environ["HF_HOME"] = "/tmp/.cache/huggingface"
+os.makedirs("/tmp/.cache/huggingface/transformers", exist_ok=True)
+
+# =============================
+# Imports
+# =============================
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import pipeline
 
+# =============================
+# Logging
+# =============================
 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. 
@@ -38,14 +52,14 @@ Please consult a dermatologist or qualified healthcare provider for personalized
 """,
 }
 
-# =========================
-# REQUEST/RESPONSE MODELS
-# =========================
+# =============================
+# FastAPI setup
+# =============================
 class ChatRequest(BaseModel):
     question: str
     context: Optional[str] = None
-    mode: Optional[str] = "dermatology"  # dermatology | general
-    max_new_tokens: Optional[int] = 200
+    mode: Optional[str] = "dermatology"  # "dermatology" | "general"
+    max_new_tokens: Optional[int] = 100
     temperature: Optional[float] = 0.7
     top_p: Optional[float] = 0.9
 
@@ -56,30 +70,35 @@ class ChatResponse(BaseModel):
     confidence: int
     sources: list
 
-# =========================
-# FASTAPI SETUP
-# =========================
 app = FastAPI(title="BioGPT-Large Medical Chatbot")
 
-MODEL_ID = "microsoft/BioGPT-Large"
+MODEL_ID = os.environ.get("MODEL_ID", "microsoft/BioGPT-Large")
 generator = None
 
+# =============================
+# Load model on startup
+# =============================
 @app.on_event("startup")
 def load_model():
     global generator
-    logger.info(f"Loading Hugging Face model via pipeline: {MODEL_ID}")
     try:
-        # Use HF hosted model (CPU is fine, HF handles backend)
+        logger.info(f"Loading Hugging Face model via pipeline: {MODEL_ID}")
         generator = pipeline("text-generation", model=MODEL_ID, device=-1)
         logger.info("Model loaded successfully.")
     except Exception as e:
         logger.exception("Failed to load model")
         generator = None
 
+# =============================
+# Root endpoint
+# =============================
 @app.get("/")
 def root():
     return {"status": "ok", "model_loaded": generator is not None, "model": MODEL_ID}
 
+# =============================
+# Chat endpoint
+# =============================
 @app.post("/chat", response_model=ChatResponse)
 def chat(req: ChatRequest):
     if generator is None:
@@ -88,14 +107,18 @@ def chat(req: ChatRequest):
     if not req.question.strip():
         raise HTTPException(status_code=400, detail="Question cannot be empty")
 
-    # Build prompt
+    # 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
 
+    logger.info(f"Generating answer for question: {req.question[:80]}...")
     t0 = time.time()
+
     try:
         outputs = generator(
             prompt,
@@ -106,8 +129,10 @@ def chat(req: ChatRequest):
             return_full_text=False,
             num_return_sequences=1,
         )
+
         answer = outputs[0]["generated_text"].strip()
         final_answer = f"{answer}\n\n{MEDICAL_PROMPTS['disclaimer']}"
+
         took = time.time() - t0
         confidence = min(95, 70 + int(len(answer) / 50))