Update app.py

app.py

@@ -234,9 +234,13 @@ def apply_scaling_or_stats(raw_vec: np.ndarray) -> (np.ndarray, Dict[str, float]
 # --------- SHAP: model wrapper & explainer ---------
 def model_proba_from_z(z_batch_np: np.ndarray) -> np.ndarray:
     """
-    Takes (N, n_features) OR a single 1D sample in z-space
-    and returns (N, K) probabilities.
-    Safe for both normal /predict and SHAP calls.
+    Batch-safe wrapper for SHAP and other callers.
+
+    Input:
+        z_batch_np: (N, n_features) or (n_features,) in z-space
+
+    Output:
+        probs: (N, K) matrix of class probabilities
     """
     z = np.array(z_batch_np, dtype=np.float32)
 
@@ -244,15 +248,31 @@ def model_proba_from_z(z_batch_np: np.ndarray) -> np.ndarray:
     if z.ndim == 1:
         z = z.reshape(1, -1)
 
-    raw = model.predict(z, verbose=0)
-    probs, _ = decode_logits(raw)
+    raw = model.predict(z, verbose=0)  # shape: (N, M)
+    if raw.ndim != 2:
+        raise ValueError(f"Unexpected raw shape from model: {raw.shape}")
 
-    # decode_logits may return (K,) if N=1, so enforce 2D
-    probs = np.array(probs, dtype=np.float32)
-    if probs.ndim == 1:
-        probs = probs.reshape(1, -1)
+    N, M = raw.shape
+    K = len(CLASSES)
 
-    return probs   # shape: (N, K)
+    if M == K - 1:
+        # CORAL: logits for K-1 thresholds → K probabilities
+        probs = coral_probs_from_logits(raw)  # (N, K)
+    elif M == K:
+        # Softmax or unnormalized scores, per row
+        exps = np.exp(raw - np.max(raw, axis=1, keepdims=True))
+        probs = exps / np.sum(exps, axis=1, keepdims=True)  # (N, K)
+    else:
+        # Fallback: row-wise normalization
+        s = np.sum(np.abs(raw), axis=1, keepdims=True)      # (N, 1)
+        probs = np.divide(
+            raw,
+            s,
+            out=np.ones_like(raw) / max(M, 1),
+            where=(s > 0),
+        )  # (N, M)
+
+    return probs
 
 
 EXPLAINER = None