Update app.py

app.py

@@ -234,23 +234,25 @@ 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:
     """
-    Wrapper for SHAP: takes (N, n_features) in z-space and returns (N, K) probabilities.
+    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.
     """
-    raw = model.predict(z_batch_np, verbose=0)
-    if raw.ndim != 2:
-        raise ValueError(f"Unexpected raw shape from model: {raw.shape}")
-    N, M = raw.shape
-    K = len(CLASSES)
+    z = np.array(z_batch_np, dtype=np.float32)
 
-    if M == K - 1:
-        probs = coral_probs_from_logits(raw)  # (N, K)
-    elif M == K:
-        exps = np.exp(raw - np.max(raw, axis=1, keepdims=True))
-        probs = exps / np.sum(exps, axis=1, keepdims=True)
-    else:
-        s = np.sum(np.abs(raw), axis=1, keepdims=True)
-        probs = np.divide(raw, s, out=np.ones_like(raw) / max(M, 1), where=(s > 0))
-    return probs
+    # Ensure 2D: (N, D)
+    if z.ndim == 1:
+        z = z.reshape(1, -1)
+
+    raw = model.predict(z, verbose=0)
+    probs, _ = decode_logits(raw)
+
+    # 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)
+
+    return probs   # shape: (N, K)
 
 
 EXPLAINER = None
@@ -341,7 +343,7 @@ async def predict(req: Request):
     Missing features are imputed if imputer present; else filled with means (if stats) or 0.
 
     This endpoint ALSO computes SHAP values for the *predicted class only*,
-    returning one SHAP value per feature (21 in total).
+    returning one SHAP value per feature (21 in total) when SHAP is available.
     """
     try:
         payload = await req.json()
@@ -367,50 +369,50 @@ async def predict(req: Request):
         shap_payload: Dict[str, Any]
 
         if not SHAP_AVAILABLE:
-            # shap library not installed in this environment
             shap_payload = {
                 "available": False,
                 "reason": "SHAP library not installed in this environment.",
             }
         else:
             try:
-                # Helper: probability function in *z-space*
-                def model_proba_from_z(z_batch_np: np.ndarray) -> np.ndarray:
-                    """
-                    Takes (N, n_features) in z-space and returns (N, K) probabilities.
-                    This mirrors the normal predict pipeline but assumes we're already in z-space.
-                    """
-                    raw_local = model.predict(z_batch_np, verbose=0)
-                    return decode_logits(raw_local)[0].reshape(-1, len(CLASSES))
-
                 # Scalar function: probability of the *predicted* class only
                 def f_scalar(z_batch):
-                    z_batch = np.array(z_batch, dtype=np.float32)
+                    """
+                    z_batch: (N, D) or (D,)
+                    returns: (N,) probability of the predicted class
+                    """
                     probs_batch = model_proba_from_z(z_batch)   # (N, K)
                     return probs_batch[:, pred_idx]             # (N,)
 
                 # Background: 50 "average" institutions at z=0
                 background_z = np.zeros((50, len(FEATURES)), dtype=np.float32)
 
-                # Create a per-call KernelExplainer for this scalar output
+                # KernelExplainer for a scalar-output model
                 explainer = shap.KernelExplainer(f_scalar, background_z)
 
-                # SHAP for this *one* observation (in z-space)
+                # SHAP for this one observation (in z-space)
                 shap_vals = explainer.shap_values(X_z, nsamples=50)
-                shap_arr = np.array(shap_vals)
-
-                # We expect shape (1, n_features) or (n_features,)
-                if shap_arr.ndim == 2 and shap_arr.shape[0] == 1:
-                    shap_vec = shap_arr[0]
+                # For scalar output, shap_vals is usually a 2D array (N, D),
+                # but some versions wrap it in a list. Handle both:
+                if isinstance(shap_vals, list):
+                    shap_mat = np.array(shap_vals[0])
                 else:
-                    shap_vec = shap_arr.reshape(-1)
+                    shap_mat = np.array(shap_vals)
 
-                if shap_vec.size != len(FEATURES):
+                # Expect (1, n_features)
+                if shap_mat.ndim == 1:
+                    shap_mat = shap_mat.reshape(1, -1)
+
+                if shap_mat.shape[0] != 1:
+                    raise ValueError(f"Unexpected SHAP batch size {shap_mat.shape[0]} (expected 1)")
+                if shap_mat.shape[1] != len(FEATURES):
                     raise ValueError(
-                        f"Unexpected SHAP vector length {shap_vec.size} "
+                        f"Unexpected SHAP vector length {shap_mat.shape[1]} "
                         f"(expected {len(FEATURES)})"
                     )
 
+                shap_vec = shap_mat[0]  # (n_features,)
+
                 shap_feature_contribs = {
                     FEATURES[i]: float(shap_vec[i]) for i in range(len(FEATURES))
                 }
@@ -452,5 +454,4 @@ async def predict(req: Request):
         return JSONResponse(
             status_code=500,
             content={"error": str(e), "trace": traceback.format_exc()},
-        )
-        
+        )