Update app.py

app.py

@@ -372,7 +372,7 @@ 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])]
 
-        resp: Dict[str, Any] = {
+                return {
             "input_ok": (len(missing) == 0),
             "missing": missing,
             "preprocess": {
@@ -380,35 +380,67 @@ async def predict(req: Request):
                 "scaler": bool(scaler),
                 "z_mode": z_mode,
             },
-            "z_scores": z_detail,  # per feature
+            "z_scores": z_detail,
             "probabilities": probs_dict,
             "predicted_state": CLASSES[pred_idx],
+            "shap": shap_out,
             "debug": {
                 "raw_shape": list(raw_logits.shape),
                 "decode_mode": mode,
                 "raw_first_row": [float(v) for v in raw_logits[0]],
             },
         }
-
+        
+        pred_idx = int(np.argmax(probs))
+        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 for predicted class ----
+                # -------- SHAP EXPLANATION (predicted class only) --------
+        shap_out = None
         if EXPLAINER is not None:
             try:
-                shap_vals_list = EXPLAINER.shap_values(X, nsamples="auto")
-                # shap_vals_list is a list of length K (classes)
-                if isinstance(shap_vals_list, list) and len(shap_vals_list) == len(CLASSES):
-                    shap_for_pred = shap_vals_list[pred_idx][0]  # (n_features,)
-                    resp["shap_target"] = CLASSES[pred_idx]
-                    resp["shap_values"] = {
-                        FEATURES[i]: float(shap_for_pred[i]) for i in range(len(FEATURES))
-                    }
+                # 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)
+                if isinstance(shap_vals, list):
+                    shap_vec = np.array(shap_vals[pred_idx][0], dtype=float)
+                    # expected_value may also be a list per class
+                    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])
+                    else:
+                        exp_val = float(exp_val_raw)
                 else:
-                    resp["shap_error"] = "Unexpected SHAP output shape."
+                    raise TypeError(f"Unsupported SHAP return type: {type(shap_vals)}")
+
+                # Map feature -> SHAP contribution (for the predicted class)
+                shap_feature_contribs = {
+                    FEATURES[i]: float(shap_vec[i])
+                    for i in range(len(FEATURES))
+                }
+
+                shap_out = {
+                    "explained_class": CLASSES[pred_idx],
+                    "expected_value": exp_val,
+                    "shap_values": shap_feature_contribs,
+                }
             except Exception as e:
-                resp["shap_error"] = f"SHAP computation failed: {repr(e)}"
+                shap_out = {
+                    "error": str(e),
+                    "trace": traceback.format_exc()
+                }
         else:
-            resp["shap_error"] = "SHAP not available in this environment."
-
-        return resp
+            shap_out = {"error": "SHAP not available on server"}
 
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e), "trace": traceback.format_exc()})