Update app.py

app.py

@@ -394,6 +394,7 @@ async def predict(req: Request):
                         if k >= len(shap_vals):
                             continue
                         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, :]
@@ -424,8 +425,29 @@ async def predict(req: Request):
                 else:
                     arr = np.array(shap_vals, dtype=float)
 
-                    # (1, K, D)
+                    # (1, D, K)  <-- THIS IS YOUR (1, 21, 5) CASE
                     if (
+                        arr.ndim == 3
+                        and arr.shape[0] == 1
+                        and arr.shape[1] == len(FEATURES)
+                        and arr.shape[2] == len(CLASSES)
+                    ):
+                        # first sample, loop over classes on last axis
+                        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,
+                        }
+
+                    # (1, K, D)
+                    elif (
                         arr.ndim == 3
                         and arr.shape[0] == 1
                         and arr.shape[1] == len(CLASSES)
@@ -436,6 +458,7 @@ async def predict(req: Request):
                             all_classes[class_name] = {
                                 FEATURES[i]: float(vec[i]) for i in range(len(FEATURES))
                             }
+
                         shap_block = {
                             "available": True,
                             "mode": "per_class",
@@ -454,6 +477,7 @@ async def predict(req: Request):
                             all_classes[class_name] = {
                                 FEATURES[i]: float(vec[i]) for i in range(len(FEATURES))
                             }
+
                         shap_block = {
                             "available": True,
                             "mode": "per_class",