Update app.py

app.py

@@ -1,70 +1,106 @@
 import torch
-from fastapi import FastAPI
+import torch.nn.functional as F
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch.nn.functional as F
 
+# ==========================================
+# 1. SETUP & CONFIGURATION
+# ==========================================
 app = FastAPI()
 
-# 1. Load Model
-MODEL_PATH = "." 
-device = torch.device("cpu")
+# Define the path to the model files (Root directory)
+MODEL_PATH = "."
+device = torch.device("cpu")  # Hugging Face Spaces (Free Tier) uses CPU
 
+# MANUAL LABEL MAPPING (Safety Net)
+# Use this to fix any confusion between Red/Green results.
+# Adjust these indices if your model predicts the wrong class.
+ID2LABEL_MANUAL = {
+    0: "neutral",
+    1: "not_shirk",
+    2: "shirk"
+}
+
+# ==========================================
+# 2. LOAD MODEL
+# ==========================================
 print("Loading model...")
 try:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
     model.to(device)
     model.eval()
-    print("Model loaded successfully!")
+    print("✅ Model loaded successfully!")
 except Exception as e:
-    print(f"CRITICAL ERROR LOADING MODEL: {e}")
-
-# 2. FORCE LABEL MAPPING (The Fix)
-# We strictly define the labels here to match your training:
-# 0 -> neutral, 1 -> not_shirk, 2 -> shirk (Alphabetical Order)
-ID2LABEL = {
-    0: "neutral",
-    1: "not_shirk",
-    2: "shirk"
-}
+    print(f"❌ CRITICAL ERROR LOADING MODEL: {e}")
+    # We do not raise an error here so the app can still start and show logs,
+    # but predictions will fail if model is None.
 
+# ==========================================
+# 3. INPUT SCHEMA
+# ==========================================
 class TextRequest(BaseModel):
     text: str
 
+# ==========================================
+# 4. API ENDPOINTS
+# ==========================================
+
 @app.get("/")
 def home():
-    return {"status": "online", "message": "BanglaBERT Shirk Detector"}
+    return {"status": "online", "system": "Dockerized BanglaBERT API"}
 
 @app.post("/predict")
 def predict(request: TextRequest):
-    # Tokenize
-    inputs = tokenizer(request.text, return_tensors="pt", truncation=True, max_length=128, padding=True)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
+    try:
+        # 1. Tokenize Input
+        inputs = tokenizer(
+            request.text, 
+            return_tensors="pt", 
+            truncation=True, 
+            max_length=128, 
+            padding=True
+        )
+        inputs = {k: v.to(device) for k, v in inputs.items()}
 
-    # Predict
-    with torch.no_grad():
-        outputs = model(**inputs)
-        probs = F.softmax(outputs.logits, dim=1)
-    
-    # Get Winner Index
-    pred_idx = torch.argmax(probs, dim=1).item()
-    
-    # ✅ FORCE CORRECT LABEL NAME
-    # We ignore the model's internal config and use our manual map
-    pred_label = ID2LABEL[pred_idx]
+        # 2. Perform Inference
+        with torch.no_grad():
+            outputs = model(**inputs)
+            probs = F.softmax(outputs.logits, dim=1)
         
-    confidence = probs[0][pred_idx].item()
-    
-    # Get All Scores with Correct Names
-    scores = {}
-    for i in range(len(probs[0])):
-        label_name = ID2LABEL[i]
-        scores[label_name] = float(probs[0][i])
+        # 3. Determine Winner
+        pred_idx = torch.argmax(probs, dim=1).item()
+        confidence = probs[0][pred_idx].item()
+        
+        # 4. Resolve Label Name
+        # Priority: Try model config first, fall back to manual map if missing
+        if model.config.id2label and len(model.config.id2label) > 0:
+            # Handle potential string/int key mismatch in config
+            pred_label = model.config.id2label.get(pred_idx, model.config.id2label.get(str(pred_idx)))
+        
+        # Fallback if config is empty or failed
+        if not pred_label:
+            pred_label = ID2LABEL_MANUAL.get(pred_idx, "unknown")
+
+        # 5. Format All Scores
+        scores = {}
+        for i in range(len(probs[0])):
+            # Get label name for this index
+            if model.config.id2label:
+                lbl = model.config.id2label.get(i, model.config.id2label.get(str(i)))
+            else:
+                lbl = ID2LABEL_MANUAL.get(i, f"LABEL_{i}")
+            
+            scores[lbl] = float(probs[0][i])
+
+        return {
+            "text": request.text,
+            "label": pred_label,
+            "confidence": confidence,
+            "scores": scores
+        }
 
-    return {
-        "text": request.text,
-        "label": pred_label,
-        "confidence": confidence,
-        "scores": scores
-    }
+    except Exception as e:
+        print(f"Prediction Error: {e}")
+        return {"error": str(e)}