Update app.py

app.py

@@ -350,107 +350,91 @@ async def predict(req: Request):
     """
     Body: JSON object mapping feature -> numeric value (strings with commas/points ok).
     Missing features are imputed if imputer present; else filled with means (if stats) or 0.
+    Returns:
+      - probabilities per state
+      - predicted_state
+      - z_scores (per feature, after imputation & scaling pipeline)
+      - shap: per-class explanations if available
     """
     try:
         payload = await req.json()
         if not isinstance(payload, dict):
-            return JSONResponse(status_code=400, content={"error": "Expected JSON object"})
+            return JSONResponse(
+                status_code=400,
+                content={"error": "Expected JSON object"},
+            )
 
-        # ---------- PREPROCESSING ----------
-        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
+        # ---------- 1) Preprocess: raw -> imputed -> z ----------
+        raw_vec = build_raw_vector(payload)          # (21,) may contain NaNs
+        raw_imp = apply_imputer_if_any(raw_vec)      # impute missing
+        z_vec, z_detail, z_mode = apply_scaling_or_stats(raw_imp)
 
-        # ---------- PREDICTION ----------
+        # ---------- 2) Model prediction ----------
         X = z_vec.reshape(1, -1).astype(np.float32)
         raw_logits = model.predict(X, verbose=0)
-        probs, mode = decode_logits(raw_logits)
+        probs, decode_mode = decode_logits(raw_logits)
 
         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])]
+        missing = [f for i, f in enumerate(FEATURES) if np.isnan(raw_vec[i])]
+
+        # ---------- 3) SHAP explanations (all classes) ----------
+        shap_block: Dict[str, Any] = {"available": False}
 
-                # ---------- SHAP EXPLANATION (predicted class only) ----------
-        shap_out = {"error": "SHAP not computed"}
-        if EXPLAINER is not None:
+        if EXPLAINER is not None and SHAP_AVAILABLE:
             try:
-                shap_vals = EXPLAINER.shap_values(X, nsamples=100)
+                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]] = {}
 
-                # 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 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 (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)
+                    # 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
+                            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,
+                    }
 
                 else:
-                    # 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)})"
-                    )
-
-                # 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)
-
-                # 5) Map feature -> contribution
-                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,
-                }
+                    # single array (1, D) – treat as "predicted class only" fallback
+                    vec = np.array(shap_vals[0], dtype=float)
+                    if vec.shape[0] == len(FEATURES):
+                        shap_block = {
+                            "available": True,
+                            "mode": "single_class",
+                            "explained_class": CLASSES[pred_idx],
+                            "values": {
+                                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)})",
+                        }
 
             except Exception as e:
-                shap_out = {"error": str(e), "trace": traceback.format_exc()}
-        else:
-            shap_out = {"error": "SHAP not available on server"}
+                shap_block = {
+                    "available": False,
+                    "error": str(e),
+                    "trace": traceback.format_exc(),
+                }
 
-        # ---------- RESPONSE ----------
+        # ---------- 4) Build response ----------
         return {
             "input_ok": (len(missing) == 0),
             "missing": missing,
@@ -459,13 +443,13 @@ async def predict(req: Request):
                 "scaler": bool(scaler),
                 "z_mode": z_mode,
             },
-            "z_scores": z_detail,          # per feature (z-space)
-            "probabilities": probs_dict,   # per class
+            "z_scores": z_detail,            # per feature
+            "probabilities": probs_dict,     # per state
             "predicted_state": CLASSES[pred_idx],
-            "shap": shap_out,              # SHAP for predicted state only
+            "shap": shap_block,
             "debug": {
                 "raw_shape": list(raw_logits.shape),
-                "decode_mode": mode,
+                "decode_mode": decode_mode,
                 "raw_first_row": [float(v) for v in raw_logits[0]],
             },
         }
@@ -473,5 +457,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()},
         )