Update app.py

app.py

@@ -2,16 +2,23 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import gradio as gr
 
 # ----------------------------------
-# Load Model
+# Initialize FastAPI
+# ----------------------------------
+app = FastAPI(title="Medical Symptom Prediction API")
+
+# ----------------------------------
+# Load Model from Hugging Face Hub
 # ----------------------------------
 MODEL_NAME = "sm89/Symptom2Disease"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-model.eval()
+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}")
 
 # ----------------------------------
 # Label Mapping
@@ -29,86 +36,57 @@ id_to_label = {
 }
 
 # ----------------------------------
-# Core Prediction Logic
+# Request Schema
 # ----------------------------------
-def predict_logic(text: str):
+class PredictionRequest(BaseModel):
+    text: str
 
-    if not text.strip():
-        raise ValueError("Text input cannot be empty")
+# ----------------------------------
+# Health Check Endpoint
+# ----------------------------------
+@app.get("/")
+def health_check():
+    return {"message": "Medical Symptom API Running"}
 
-    inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        truncation=True,
-        padding=True,
-        max_length=128
-    )
+# ----------------------------------
+# Prediction Endpoint
+# ----------------------------------
+@app.post("/predict")
+def predict(request: PredictionRequest):
 
-    with torch.no_grad():
-        outputs = model(**inputs)
-        probabilities = torch.softmax(outputs.logits, dim=1)
+    if not request.text.strip():
+        raise HTTPException(status_code=400, detail="Text input cannot be empty")
 
-    top_probs, top_indices = torch.topk(probabilities, 3)
+    try:
+        inputs = tokenizer(
+            request.text,
+            return_tensors="pt",
+            truncation=True,
+            padding=True,
+            max_length=128
+        )
 
-    results = []
+        with torch.no_grad():
+            outputs = model(**inputs)
+            probabilities = torch.softmax(outputs.logits, dim=1)
 
-    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)
-        })
+        top_probs, top_indices = torch.topk(probabilities, 3)
 
-    return {
-        "input_text": text,
-        "top_predictions": results,
-        "final_prediction": results[0]
-    }
+        results = []
 
-# ----------------------------------
-# FastAPI App
-# ----------------------------------
-app = FastAPI(title="Medical Symptom Prediction API")
+        for prob, idx in zip(top_probs[0], top_indices[0]):
+            label_index = int(idx.item())
 
-class PredictionRequest(BaseModel):
-    text: str
+            results.append({
+                "department": id_to_label.get(label_index, f"LABEL_{label_index}"),
+                "confidence": round(float(prob.item()), 4)
+            })
 
-@app.get("/health")
-def health():
-    return {"status": "running"}
+        return {
+            "input_text": request.text,
+            "top_predictions": results,
+            "final_prediction": results[0]
+        }
 
-@app.post("/predict")
-def predict(request: PredictionRequest):
-    try:
-        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
-# ----------------------------------
-def gradio_predict(text):
-    try:
-        result = predict_logic(text)
-        output = ""
-
-        for item in result["top_predictions"]:
-            output += f"{item['department']} ({item['confidence']})\n"
-
-        return output
-
-    except Exception as e:
-        return str(e)
-
-demo = gr.Interface(
-    fn=gradio_predict,
-    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 (IMPORTANT)
-app = gr.mount_gradio_app(app, demo, path="/ui")