Update app.py

app.py

@@ -1,24 +1,17 @@
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import gradio as gr
 
 # ----------------------------------
-# Initialize FastAPI
-# ----------------------------------
-app = FastAPI(title="Medical Symptom Prediction API")
-
-# ----------------------------------
-# Load Model from Hugging Face Hub
+# Load Model
 # ----------------------------------
 MODEL_NAME = "sm89/Symptom2Disease"
 
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-    model.eval()
-except Exception as e:
-    raise RuntimeError(f"Model loading failed: {e}")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+model.eval()
 
 # ----------------------------------
 # Label Mapping
@@ -61,7 +54,6 @@ def predict_logic(text: str):
 
     for prob, idx in zip(top_probs[0], top_indices[0]):
         label_index = int(idx.item())
-
         results.append({
             "department": id_to_label.get(label_index, f"LABEL_{label_index}"),
             "confidence": round(float(prob.item()), 4)
@@ -74,58 +66,49 @@ def predict_logic(text: str):
     }
 
 # ----------------------------------
-# Health Endpoint
+# FastAPI App
 # ----------------------------------
+app = FastAPI(title="Medical Symptom Prediction API")
+
+class PredictionRequest(BaseModel):
+    text: str
+
 @app.get("/health")
-def health_check():
-    return {"message": "Medical Symptom API Running"}
+def health():
+    return {"status": "running"}
 
-# ----------------------------------
-# JSON Prediction Endpoint (No 422 Issue)
-# ----------------------------------
 @app.post("/predict")
-async def predict_api(request: Request):
+def predict(request: PredictionRequest):
     try:
-        body = await request.json()
-        text = body.get("text", "")
-
-        if not text.strip():
-            raise HTTPException(status_code=400, detail="Text input cannot be empty")
-
-        return predict_logic(text)
-
-    except Exception as e:
+        return predict_logic(request.text)
+    except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 # ----------------------------------
-# Gradio UI Function
+# Gradio UI
 # ----------------------------------
 def gradio_predict(text):
     try:
         result = predict_logic(text)
-
-        output = "Top Predictions:\n\n"
+        output = ""
 
         for item in result["top_predictions"]:
-            output += f"{item['department']} → {item['confidence']}\n"
-
-        output += f"\nFinal Prediction: {result['final_prediction']['department']}"
+            output += f"{item['department']} ({item['confidence']})\n"
 
         return output
 
     except Exception as e:
         return str(e)
 
-# ----------------------------------
-# Create Gradio Interface
-# ----------------------------------
 demo = gr.Interface(
     fn=gradio_predict,
-    inputs=gr.Textbox(lines=3, placeholder="Enter symptoms here..."),
-    outputs=gr.Textbox(label="Prediction Result"),
-    title="Medical Symptom to Department Predictor",
-    description="Enter symptoms and get top 3 predicted departments."
+    inputs=gr.Textbox(lines=3, placeholder="Enter symptoms here"),
+    outputs=gr.Textbox(label="Prediction"),
+    title="Medical Symptom Predictor",
+    description="Enter symptoms to get top predicted medical departments"
 )
 
-# Mount Gradio at /ui (safe way)
+# Mount Gradio at /ui (IMPORTANT)
 app = gr.mount_gradio_app(app, demo, path="/ui")