Update app.py

app.py

@@ -30,7 +30,7 @@ class SimpleStacking:
 
         for model in self.base_models.values():
             if hasattr(model, "predict_proba"):
-                probs = model.predict_proba(X)[:, 1]
+                probs = model.predict_proba(X)
                 meta_features.append(probs)
             else:
                 preds = model.predict(X)
@@ -42,7 +42,8 @@ class SimpleStacking:
             return self.final_model.predict_proba(meta_X)
         else:
             preds = self.final_model.predict(meta_X)
-            return np.column_stack([1 - preds, preds])
+            # For simplicity, assume predict gives class, but need probs
+            return preds
 
 
 # ======================
@@ -90,23 +91,22 @@ def predict_potability(ph, tds, turbidity, temperature):
         "Stacking": stacking_model
     }
 
+    class_names = ["Non-Potable", "Borderline", "Potable"]
+
     outputs = {}
 
     for name, model in models.items():
         pred = model.predict(scaled)[0]
 
         if hasattr(model, "predict_proba"):
-            prob = model.predict_proba(scaled)[0][1]
-            prob_text = f"{prob:.4f}"
+            probs = model.predict_proba(scaled)[0]
+            prob_str = ", ".join([f"{class_names[i]}: {probs[i]:.4f}" for i in range(3)])
         else:
-            prob_text = "N/A"
+            prob_str = "N/A"
 
-        label = "Potable" if pred == 1 else "Non-Potable"
+        label = class_names[pred]
 
-        if prob_text != "N/A":
-            outputs[name] = f"{label} (Probability: {prob_text})"
-        else:
-            outputs[name] = label
+        outputs[name] = f"{label} (Probabilities: {prob_str})"
 
     return (
         outputs["XGBoost"],