Update app.py

app.py

@@ -370,50 +370,70 @@ 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[i])]
 
-        # ---------- SHAP EXPLANATION (predicted class only) ----------
+                # ---------- SHAP EXPLANATION (predicted class only) ----------
         shap_out = {"error": "SHAP not computed"}
         if EXPLAINER is not None:
             try:
                 shap_vals = EXPLAINER.shap_values(X, nsamples=100)
 
-                # 1) Pull out the array for the predicted class (if multi-output)
+                # 1) Pull raw SHAP tensor
                 if isinstance(shap_vals, list):
+                    # Classic multi-output: list[len = n_classes], each (n_samples, n_features)
                     raw_sv = np.array(shap_vals[pred_idx])
                 else:
+                    # Single array, possibly (n_samples, n_features) or (n_samples, n_features, n_outputs)
                     raw_sv = np.array(shap_vals)
 
-                # 2) Normalize shapes: we want a 1D vector of length n_features
-                # Possible shapes we might see:
-                #   (n_features,)
-                #   (1, n_features)
-                #   (n_samples, n_features) -> take first sample
+                # 2) Normalize shapes to a 1D vector (n_features,) for the predicted class
                 if raw_sv.ndim == 1:
+                    # Already (n_features,)
                     shap_vec = raw_sv.astype(float)
+
                 elif raw_sv.ndim == 2:
+                    # (n_samples, n_features) or (n_features, 1)
                     if raw_sv.shape[0] == 1:
+                        # (1, n_features)
                         shap_vec = raw_sv[0].astype(float)
+                    elif raw_sv.shape[1] == 1:
+                        # (n_features, 1)
+                        shap_vec = raw_sv[:, 0].astype(float)
                     else:
-                        # assume shape (n_samples, n_features); take first sample
+                        # assume (n_samples, n_features), take first sample
                         shap_vec = raw_sv[0].astype(float)
+
+                elif raw_sv.ndim == 3:
+                    # Most likely (n_samples, n_features, n_outputs)
+                    n_samples, n_features, n_outputs = raw_sv.shape
+                    if n_samples < 1:
+                        raise ValueError(f"SHAP 3D output has zero samples: {raw_sv.shape}")
+                    if pred_idx >= n_outputs:
+                        raise ValueError(
+                            f"SHAP 3D output has only {n_outputs} outputs, "
+                            f"cannot index class {pred_idx}"
+                        )
+                    # take first sample, all features, predicted class
+                    shap_vec = raw_sv[0, :, pred_idx].astype(float)
+
                 else:
-                    # last resort: flatten sample dims, take first "row"
-                    raw_sv = raw_sv.reshape(raw_sv.shape[0], -1)
-                    shap_vec = raw_sv[0].astype(float)
+                    # Fallback: flatten all sample dims, keep first feature-block
+                    flat = raw_sv.reshape(raw_sv.shape[0], -1)
+                    shap_vec = flat[0].astype(float)
 
+                # 3) Sanity check length
                 if shap_vec.shape[0] != len(FEATURES):
                     raise ValueError(
                         f"Unexpected SHAP vector length {shap_vec.shape[0]} "
                         f"(expected {len(FEATURES)})"
                     )
 
-                # 3) Expected value: baseline logit/prob for that class
+                # 4) Expected value (baseline) for the predicted class
                 exp_raw = EXPLAINER.expected_value
                 if isinstance(exp_raw, (list, np.ndarray)):
                     exp_val = float(np.array(exp_raw)[pred_idx])
                 else:
                     exp_val = float(exp_raw)
 
-                # 4) Map feature -> SHAP contribution
+                # 5) Map feature -> contribution
                 shap_feature_contribs = {
                     FEATURES[i]: float(shap_vec[i])
                     for i in range(len(FEATURES))