Update app.py

app.py

@@ -378,41 +378,103 @@ async def predict(req: Request):
         probs_dict = {CLASSES[i]: float(probs[i]) for i in range(len(CLASSES))}
         missing = [f for i, f in enumerate(FEATURES) if np.isnan(raw_vec[i])]
 
-        # ---------- 3) SHAP explanations (all classes) ----------
+               # ---------- 3) SHAP explanations (all classes) ----------
         shap_block: Dict[str, Any] = {"available": False}
 
         if EXPLAINER is not None and SHAP_AVAILABLE:
             try:
                 X_z = z_vec.reshape(1, -1).astype(np.float32)
-                # KernelExplainer: usually returns list of length K (one (1,D) array per class)
                 shap_vals = EXPLAINER.shap_values(X_z, nsamples=50)
 
                 all_classes: Dict[str, Dict[str, float]] = {}
 
+                # ---------- CASE 1: SHAP returns list (usual multi-class) ----------
                 if isinstance(shap_vals, list):
-                    # per-class outputs
                     for k, class_name in enumerate(CLASSES):
                         if k >= len(shap_vals):
                             continue
-                        vec = np.array(shap_vals[k][0], dtype=float)  # shape (D,)
-                        if vec.shape[0] != len(FEATURES):
-                            # shape mismatch: bail out gracefully
+                        arr = np.array(shap_vals[k], dtype=float)  # shape (N, D) or (D,)
+                        # reduce to a 1D (D,) vector for the first sample
+                        if arr.ndim == 2 and arr.shape[0] >= 1 and arr.shape[1] == len(FEATURES):
+                            vec = arr[0, :]
+                        elif arr.ndim == 1 and arr.shape[0] == len(FEATURES):
+                            vec = arr
+                        else:
+                            # shape we don't know how to handle for this class
                             continue
+
                         all_classes[class_name] = {
                             FEATURES[i]: float(vec[i]) for i in range(len(FEATURES))
                         }
 
-                    shap_block = {
-                        "available": True,
-                        "mode": "per_class",
-                        "explained_classes": list(all_classes.keys()),
-                        "all_classes": all_classes,
-                    }
+                    if all_classes:
+                        shap_block = {
+                            "available": True,
+                            "mode": "per_class",
+                            "explained_classes": list(all_classes.keys()),
+                            "all_classes": all_classes,
+                        }
+                    else:
+                        shap_block = {
+                            "available": False,
+                            "error": "No per-class SHAP vectors matched expected shape.",
+                        }
 
+                # ---------- CASE 2: SHAP returns a numpy array ----------
                 else:
-                    # single array (1, D) – treat as "predicted class only" fallback
-                    vec = np.array(shap_vals[0], dtype=float)
-                    if vec.shape[0] == len(FEATURES):
+                    arr = np.array(shap_vals, dtype=float)
+
+                    # (1, K, D)
+                    if (
+                        arr.ndim == 3
+                        and arr.shape[0] == 1
+                        and arr.shape[1] == len(CLASSES)
+                        and arr.shape[2] == len(FEATURES)
+                    ):
+                        for k, class_name in enumerate(CLASSES):
+                            vec = arr[0, k, :]  # (D,)
+                            all_classes[class_name] = {
+                                FEATURES[i]: float(vec[i]) for i in range(len(FEATURES))
+                            }
+                        shap_block = {
+                            "available": True,
+                            "mode": "per_class",
+                            "explained_classes": list(all_classes.keys()),
+                            "all_classes": all_classes,
+                        }
+
+                    # (K, D)
+                    elif (
+                        arr.ndim == 2
+                        and arr.shape[0] == len(CLASSES)
+                        and arr.shape[1] == len(FEATURES)
+                    ):
+                        for k, class_name in enumerate(CLASSES):
+                            vec = arr[k, :]  # (D,)
+                            all_classes[class_name] = {
+                                FEATURES[i]: float(vec[i]) for i in range(len(FEATURES))
+                            }
+                        shap_block = {
+                            "available": True,
+                            "mode": "per_class",
+                            "explained_classes": list(all_classes.keys()),
+                            "all_classes": all_classes,
+                        }
+
+                    # Single-vector fallback: (1, D) or (D,)
+                    elif arr.ndim == 2 and arr.shape[0] == 1 and arr.shape[1] == len(FEATURES):
+                        vec = arr[0, :]  # (D,)
+                        shap_block = {
+                            "available": True,
+                            "mode": "single_class",
+                            "explained_class": CLASSES[pred_idx],
+                            "values": {
+                                FEATURES[i]: float(vec[i]) for i in range(len(FEATURES))
+                            },
+                        }
+
+                    elif arr.ndim == 1 and arr.shape[0] == len(FEATURES):
+                        vec = arr  # (D,)
                         shap_block = {
                             "available": True,
                             "mode": "single_class",
@@ -421,10 +483,11 @@ async def predict(req: Request):
                                 FEATURES[i]: float(vec[i]) for i in range(len(FEATURES))
                             },
                         }
+
                     else:
                         shap_block = {
                             "available": False,
-                            "error": f"Unexpected SHAP vector length {vec.shape[0]} (expected {len(FEATURES)})",
+                            "error": f"Unexpected SHAP array shape {arr.shape}",
                         }
 
             except Exception as e:
@@ -433,7 +496,6 @@ async def predict(req: Request):
                     "error": str(e),
                     "trace": traceback.format_exc(),
                 }
-
         # ---------- 4) Build response ----------
         return {
             "input_ok": (len(missing) == 0),