Update app.py

app.py

@@ -34,26 +34,63 @@ def load_pipeline():
             print(f"Failed to load {model_id}: {e}")
     raise RuntimeError("No sentiment model could be loaded")
 
-def scores_to_label(scores, binary=False):
-    m = {s["label"].lower(): s["score"] for s in scores}
-    if binary or ("neutral" not in m):
-        neg, pos = m.get("negative", 0.0), m.get("positive", 0.0)
-        score = pos - neg  # [-1,1]
-        conf = max(neg, pos)
+def scores_to_label(scores, binary=None):
+    m = {s["label"].lower(): float(s["score"]) for s in scores}
+    keys = set(m.keys())
+
+    neg = neu = pos = 0.0
+    detected_binary = False
+
+    if {"negative", "positive"} & keys:
+        # Named labels present; detect if neutral exists
+        neg = m.get("negative", 0.0)
+        pos = m.get("positive", 0.0)
+        neu = m.get("neutral", 0.0)
+        detected_binary = ("neutral" not in m) or (len(m) == 2)
+    elif any(k.startswith("label_") for k in keys):
+        neg = m.get("label_0", 0.0)
+        if "label_2" in m or len(m) >= 3:
+            neu = m.get("label_1", 0.0)
+            pos = m.get("label_2", 0.0)
+            detected_binary = False
+        else:
+            pos = m.get("label_1", 0.0)
+            neu = 0.0
+            detected_binary = True
+    else:
+        for k, v in m.items():
+            if "pos" in k: pos = v
+            if "neg" in k: neg = v
+            if "neu" in k: neu = v
+        detected_binary = (neu == 0.0)
+
+    is_binary = detected_binary if binary is None else bool(binary)
+
+    if is_binary:
+        score = pos - neg
+        conf = max(pos, neg)
         label = "positive" if score > 0 else ("negative" if score < 0 else "neutral")
-        return {"label": label, "score": max(-1.0, min(1.0, score)), "confidence": conf,
-                "scores": {"positive": pos, "neutral": 0.0, "negative": neg}}
+        return {
+            "label": label,
+            "score": max(-1.0, min(1.0, score)),
+            "confidence": conf,
+            "scores": {"positive": pos, "neutral": 0.0, "negative": neg},
+        }
     else:
-        neg, neu, pos = m.get("negative", 0.0), m.get("neutral", 0.0), m.get("positive", 0.0)
         denom = max(1e-6, 1.0 - neu)
-        score = (pos - neg) / denom  
+        score = (pos - neg) / denom
         score = max(-1.0, min(1.0, score))
-        conf = max(neg, neu, pos)
+        conf = max(pos, neg, neu)
         label = "positive" if pos >= max(neg, neu) else ("negative" if neg >= max(pos, neu) else "neutral")
         if conf < 0.55:
             label = "neutral"
-        return {"label": label, "score": score, "confidence": conf,
-                "scores": {"positive": pos, "neutral": neu, "negative": neg}}
+        return {
+            "label": label,
+            "score": score,
+            "confidence": conf,
+            "scores": {"positive": pos, "neutral": neu, "negative": neg},
+        }
+
 
 @app.post("/predict")
 def predict(payload: Payload):