Update app.py

app.py

@@ -350,20 +350,22 @@ 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.
-
-    Now also returns SHAP values for the predicted_state (if SHAP is available).
     """
     try:
         payload = await req.json()
         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
 
-        # Predict
+        # ---------- PREDICTION ----------
         X = z_vec.reshape(1, -1).astype(np.float32)
         raw_logits = model.predict(X, verbose=0)
         probs, mode = decode_logits(raw_logits)
@@ -372,32 +374,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])]
 
-        return {
-            "input_ok": (len(missing) == 0),
-            "missing": missing,
-            "preprocess": {
-                "imputer": bool(imputer),
-                "scaler": bool(scaler),
-                "z_mode": z_mode,
-            },
-            "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
+        # ---------- SHAP EXPLANATION (predicted class only) ----------
         if EXPLAINER is not None:
             try:
                 # X is already z-space: shape (1, n_features)
@@ -406,12 +383,12 @@ async def predict(req: Request):
                 # 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)
@@ -420,6 +397,8 @@ async def predict(req: Request):
                         exp_val = float(exp_val_raw[0])
                     else:
                         exp_val = float(exp_val_raw)
+
+                # Anything else – we consider wrong type
                 else:
                     raise TypeError(f"Unsupported SHAP return type: {type(shap_vals)}")
 
@@ -434,6 +413,7 @@ async def predict(req: Request):
                     "expected_value": exp_val,
                     "shap_values": shap_feature_contribs,
                 }
+
             except Exception as e:
                 shap_out = {
                     "error": str(e),
@@ -442,5 +422,31 @@ async def predict(req: Request):
         else:
             shap_out = {"error": "SHAP not available on server"}
 
+        # ---------- RESPONSE ----------
+        return {
+            "input_ok": (len(missing) == 0),
+            "missing": missing,
+            "preprocess": {
+                "imputer": bool(imputer),
+                "scaler": bool(scaler),
+                "z_mode": z_mode,
+            },
+            "z_scores": z_detail,          # per feature
+            "probabilities": probs_dict,   # per class
+            "predicted_state": CLASSES[pred_idx],
+            "shap": shap_out,              # SHAP for predicted state only
+            "debug": {
+                "raw_shape": list(raw_logits.shape),
+                "decode_mode": mode,
+                "raw_first_row": [float(v) for v in raw_logits[0]],
+            },
+        }
+
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e), "trace": traceback.format_exc()})
+        return JSONResponse(
+            status_code=500,
+            content={
+                "error": str(e),
+                "trace": traceback.format_exc()
+            }
+        )