Update app.py

app.py

@@ -356,11 +356,7 @@ async def predict(req: Request):
         if not isinstance(payload, dict):
             return JSONResponse(status_code=400, content={"error": "Expected JSON object"})
 
-        # default SHAP block – will be overwritten if explanation succeeds
-        shap_out = {"error": "SHAP not computed"}
-
         # ---------- PREPROCESSING ----------
-        # Build in EXACT training order
         raw = build_raw_vector(payload)            # may contain NaNs
         raw_imp = apply_imputer_if_any(raw)        # impute
         z_vec, z_detail, z_mode = apply_scaling_or_stats(raw_imp)  # scale / z-score
@@ -375,34 +371,49 @@ async def predict(req: Request):
         missing = [f for i, f in enumerate(FEATURES) if np.isnan(raw[i])]
 
         # ---------- SHAP EXPLANATION (predicted class only) ----------
+        shap_out = {"error": "SHAP not computed"}
         if EXPLAINER is not None:
             try:
-                # X is already z-space: shape (1, n_features)
                 shap_vals = EXPLAINER.shap_values(X, nsamples=100)
 
-                # Case 1: multi-output -> list of length K, each (1, n_features)
+                # 1) Pull out the array for the predicted class (if multi-output)
                 if isinstance(shap_vals, list):
-                    shap_vec = np.array(shap_vals[pred_idx][0], dtype=float)
-                    exp_val_raw = EXPLAINER.expected_value
-                    if isinstance(exp_val_raw, (list, np.ndarray)):
-                        exp_val = float(exp_val_raw[pred_idx])
-                    else:
-                        exp_val = float(exp_val_raw)
-
-                # Case 2: single-output -> ndarray (1, n_features)
-                elif isinstance(shap_vals, np.ndarray):
-                    shap_vec = np.array(shap_vals[0], dtype=float)
-                    exp_val_raw = EXPLAINER.expected_value
-                    if isinstance(exp_val_raw, (list, np.ndarray)):
-                        exp_val = float(exp_val_raw[0])
+                    raw_sv = np.array(shap_vals[pred_idx])
+                else:
+                    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
+                if raw_sv.ndim == 1:
+                    shap_vec = raw_sv.astype(float)
+                elif raw_sv.ndim == 2:
+                    if raw_sv.shape[0] == 1:
+                        shap_vec = raw_sv[0].astype(float)
                     else:
-                        exp_val = float(exp_val_raw)
-
-                # Anything else – we consider wrong type
+                        # assume shape (n_samples, n_features); take first sample
+                        shap_vec = raw_sv[0].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)
+
+                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
+                exp_raw = EXPLAINER.expected_value
+                if isinstance(exp_raw, (list, np.ndarray)):
+                    exp_val = float(np.array(exp_raw)[pred_idx])
                 else:
-                    raise TypeError(f"Unsupported SHAP return type: {type(shap_vals)}")
+                    exp_val = float(exp_raw)
 
-                # Map feature -> SHAP contribution (for the predicted class)
+                # 4) Map feature -> SHAP contribution
                 shap_feature_contribs = {
                     FEATURES[i]: float(shap_vec[i])
                     for i in range(len(FEATURES))
@@ -415,10 +426,7 @@ async def predict(req: Request):
                 }
 
             except Exception as e:
-                shap_out = {
-                    "error": str(e),
-                    "trace": traceback.format_exc()
-                }
+                shap_out = {"error": str(e), "trace": traceback.format_exc()}
         else:
             shap_out = {"error": "SHAP not available on server"}
 
@@ -431,7 +439,7 @@ async def predict(req: Request):
                 "scaler": bool(scaler),
                 "z_mode": z_mode,
             },
-            "z_scores": z_detail,          # per feature
+            "z_scores": z_detail,          # per feature (z-space)
             "probabilities": probs_dict,   # per class
             "predicted_state": CLASSES[pred_idx],
             "shap": shap_out,              # SHAP for predicted state only
@@ -445,8 +453,5 @@ async def predict(req: Request):
     except Exception as e:
         return JSONResponse(
             status_code=500,
-            content={
-                "error": str(e),
-                "trace": traceback.format_exc()
-            }
+            content={"error": str(e), "trace": traceback.format_exc()}
         )