Update main.py

main.py

@@ -11,7 +11,6 @@ import os
 
 app = FastAPI(title="MedGuard API")
 
-# --- CORS CONFIGURATION ---
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], 
@@ -20,7 +19,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- CONFIGURATION ---
 MODEL_PATH = "./model"
 DEVICE = "cpu"
 
@@ -28,20 +26,31 @@ print(f"🔄 Loading Model from {MODEL_PATH}...")
 model = None
 tokenizer = None
 
+# --- CRITICAL FIX: MATCH TRAINING LABEL MAP ---
+# Training Map: {'Not Relevant': 0, 'Partially Relevant': 1, 'Highly Relevant': 2}
+# This list MUST follow the index order: [Index 0, Index 1, Index 2]
+LABELS = ["Not Relevant", "Partially Relevant", "Highly Relevant"]
+
 try:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
     model.to(DEVICE)
     model.eval()
-    print("✅ Model Loaded Successfully!")
+    
+    # Validation check (Optional but good)
+    if model.config.id2label:
+        print(f"ℹ️ Model config labels: {model.config.id2label}")
+        # We enforce our manual list because sometimes configs get messed up during saving
+        # but you should visually verify if this print matches our LABELS list
+        
+    print(f"✅ Model Loaded! Label Mapping: {LABELS}")
+    
 except Exception as e:
     print(f"❌ Error loading local model: {e}")
-    # Fallback
     MODEL_NAME = "csebuetnlp/banglabert"
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3)
 
-# --- DATA MODELS ---
 class QueryRequest(BaseModel):
     genre: str = ""
     prompt: str = ""
@@ -53,8 +62,6 @@ class PredictionResponse(BaseModel):
     probs: dict
     explanation: list = None
 
-LABELS = ["Highly Relevant", "Partially Relevant", "Not Relevant"]
-
 def predict_proba_lime(texts):
     inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=128).to(DEVICE)
     with torch.no_grad():
@@ -63,7 +70,7 @@ def predict_proba_lime(texts):
 
 @app.get("/")
 def health_check():
-    return {"status": "active", "model": "MedGuard v2.1 (Space Concatenation)"}
+    return {"status": "active", "model": "MedGuard v2.3 (Fixed Labels)"}
 
 @app.post("/predict", response_model=PredictionResponse)
 def predict(request: QueryRequest):
@@ -71,15 +78,12 @@ def predict(request: QueryRequest):
         raise HTTPException(status_code=503, detail="Model not loaded")
     
     try:
-        # --- CRITICAL FIX ---
-        # We use simple SPACE concatenation to match standard training dataframe practices.
-        # No [SEP] tokens, just "Genre Prompt Response"
+        # Use simple space concatenation
         parts = [part for part in [request.genre, request.prompt, request.text] if part]
         full_input = " ".join(parts)
             
-        print(f"📥 Analyzing: {full_input[:100]}...") 
+        print(f"📥 Analyzing: {full_input[:50]}...") 
 
-        # 1. PREDICT
         inputs = tokenizer(full_input, return_tensors="pt", truncation=True, max_length=128).to(DEVICE)
         with torch.no_grad():
             outputs = model(**inputs)
@@ -87,7 +91,12 @@ def predict(request: QueryRequest):
         
         pred_idx = np.argmax(probs)
         
-        # 2. EXPLAIN (LIME)
+        # Ensure index is valid
+        if pred_idx >= len(LABELS):
+            label_str = "Unknown"
+        else:
+            label_str = LABELS[pred_idx]
+
         explainer = LimeTextExplainer(
             class_names=LABELS, 
             split_expression=lambda x: x.split() 
@@ -103,7 +112,7 @@ def predict(request: QueryRequest):
         lime_features = exp.as_list(label=pred_idx) 
         
         return {
-            "label": LABELS[pred_idx],
+            "label": label_str,
             "confidence": round(float(probs[pred_idx]) * 100, 2),
             "probs": {l: round(float(p), 4) for l, p in zip(LABELS, probs)},
             "explanation": lime_features