Update app.py

app.py

@@ -41,6 +41,11 @@ def _load_model():
             _ = _model(**_tokenizer(["warm up"], return_tensors="pt")).logits
 
 
+def _id2label():
+    cfg = getattr(_model, "config", None)
+    return getattr(cfg, "id2label", {}) if cfg else {}
+
+
 @app.get("/")
 def root():
     _load_model()
@@ -77,18 +82,30 @@ def predict(payload: PredictPayload):
             logits = outputs.logits  # [1, num_labels]
             logits_list = logits[0].tolist()
             pred_idx = int(torch.argmax(logits, dim=-1).item())
+
+            # Keep client-compatible fields but also provide raw outputs
+            probs_t = torch.softmax(logits, dim=-1)[0]
+            score = float(probs_t[pred_idx])
+
     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))
+    id2label = _id2label()
+    # Resolve label from model config (support int or str keys)
+    pred_label = id2label.get(pred_idx, id2label.get(str(pred_idx), str(pred_idx)))
 
+    # Build per-label probabilities for debugging/verification
+    probs = {}
+    for i, p in enumerate(probs_t.tolist()):
+        probs[id2label.get(i, id2label.get(str(i), str(i)))] = float(p)
+
+    # Backward-compatible keys: "label" and "score"
     return {
-        "logits": logits_list,           # direct model output (authoritative)
-        "predicted_index": pred_idx,     # argmax over logits
-        "predicted_label": pred_label,   # from model config if available
+        "label": pred_label,            # expected by your client
+        "score": score,                 # probability of predicted class (softmax)
+        "predicted_index": pred_idx,    # raw argmax index from logits
+        "logits": logits_list,          # raw model output
+        "probs": probs,                 # per-label probabilities
         "id2label": id2label,
-        "label2id": getattr(cfg, "label2id", {}) if cfg else {},
+        "label2id": getattr(getattr(_model, "config", None), "label2id", {}),
     }