Update app.py

app.py

@@ -35,39 +35,10 @@ def _load_model():
     if _tokenizer is None or _model is None:
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
         _model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, cache_dir=CACHE_DIR)
-        _model.eval()  # disable dropout for inference
+        _model.eval()  # inference mode
         # Warm-up
         with torch.no_grad():
             _ = _model(**_tokenizer(["warm up"], return_tensors="pt")).logits
-        # Log mapping once for verification
-        try:
-            print("id2label:", _model.config.id2label)
-            print("label2id:", _model.config.label2id)
-        except Exception:
-            pass
-
-
-def _label_for_index(idx: int) -> str:
-    cfg = getattr(_model, "config", None)
-    id2label = getattr(cfg, "id2label", None) if cfg else None
-    if isinstance(id2label, dict):
-        # Try int key first, then str key
-        if idx in id2label:
-            return id2label[idx]
-        sidx = str(idx)
-        if sidx in id2label:
-            return id2label[sidx]
-    # Fallback for 2-class default
-    default_map = {0: "LEGIT", 1: "PHISH"}
-    return default_map.get(idx, str(idx))
-
-
-def _probs_dict(probs_list):
-    # Map each index to its label name for readability
-    out = {}
-    for i, p in enumerate(probs_list):
-        out[_label_for_index(i)] = float(p)
-    return out
 
 
 @app.get("/")
@@ -85,9 +56,10 @@ def root():
 def labels():
     _load_model()
     cfg = getattr(_model, "config", None)
-    id2label = getattr(cfg, "id2label", {}) if cfg else {}
-    label2id = getattr(cfg, "label2id", {}) if cfg else {}
-    return {"id2label": id2label, "label2id": label2id}
+    return {
+        "id2label": getattr(cfg, "id2label", {}) if cfg else {},
+        "label2id": getattr(cfg, "label2id", {}) if cfg else {},
+    }
 
 
 @app.post("/predict")
@@ -101,17 +73,22 @@ def predict(payload: PredictPayload):
                 truncation=True,
                 max_length=512
             )
-            logits = _model(**inputs).logits  # [1, num_labels]
-            probs_t = torch.softmax(logits, dim=-1)[0]  # [num_labels]
-            probs = probs_t.tolist()
-            pred_idx = int(torch.argmax(probs_t).item())
-            pred_label = _label_for_index(pred_idx)
-            pred_score = float(probs[pred_idx])
+            outputs = _model(**inputs)
+            logits = outputs.logits  # [1, num_labels]
+            logits_list = logits[0].tolist()
+            pred_idx = int(torch.argmax(logits, dim=-1).item())
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})
 
+    cfg = getattr(_model, "config", None)
+    id2label = getattr(cfg, "id2label", {}) if cfg else {}
+    # Try to fetch label from config without any remapping logic
+    pred_label = id2label.get(pred_idx, id2label.get(str(pred_idx), None))
+
     return {
-        "label": pred_label,
-        "score": pred_score,
-        "probs": _probs_dict(probs),
+        "logits": logits_list,           # direct model output (authoritative)
+        "predicted_index": pred_idx,     # argmax over logits
+        "predicted_label": pred_label,   # from model config if available
+        "id2label": id2label,
+        "label2id": getattr(cfg, "label2id", {}) if cfg else {},
     }