Update app.py

app.py

@@ -7,7 +7,7 @@ from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
-
+# ---------- SHAP optional import ----------
 try:
     import shap
     SHAP_AVAILABLE = True
@@ -129,7 +129,7 @@ else:
     print("⚠️ No scaler found — using manual z-scoring if stats are available.")
 
 # Stats (means/std) for fallback manual z-score
-stats = {}
+stats: Dict[str, Dict[str, float]] = {}
 if os.path.isfile(STATS_PATH):
     stats = load_json(STATS_PATH)
     print(f"Loaded means/std from {STATS_PATH}")
@@ -147,6 +147,8 @@ def coral_probs_from_logits(logits_np: np.ndarray) -> np.ndarray:
     left  = tf.concat([tf.ones_like(sig[:, :1]), sig], axis=1)
     right = tf.concat([sig, tf.zeros_like(sig[:, :1])], axis=1)
     probs = tf.clip_by_value(left - right, 1e-12, 1.0)
+    # normalize row-wise just in case
+    probs = probs / tf.reduce_sum(probs, axis=1, keepdims=True)
     return probs.numpy()
 
 
@@ -236,8 +238,47 @@ def apply_scaling_or_stats(raw_vec: np.ndarray) -> (np.ndarray, Dict[str, float]
         return z, z_detail, "manual_stats"
 
 
+# --------- SHAP model wrapper & explainer ---------
+def model_proba_from_z(z_batch_np: np.ndarray) -> np.ndarray:
+    """
+    Wrapper for SHAP: takes (N, n_features) in z-space and returns (N, K) probabilities.
+    """
+    raw = model.predict(z_batch_np, verbose=0)
+    if raw.ndim != 2:
+        raise ValueError(f"Unexpected raw shape from model: {raw.shape}")
+    N, M = raw.shape
+    K = len(CLASSES)
+
+    if M == K - 1:
+        # CORAL
+        probs = coral_probs_from_logits(raw)  # (N, K)
+    elif M == K:
+        # Softmax or scores
+        exps = np.exp(raw - np.max(raw, axis=1, keepdims=True))
+        probs = exps / np.sum(exps, axis=1, keepdims=True)
+    else:
+        # Fallback normalize
+        s = np.sum(np.abs(raw), axis=1, keepdims=True)
+        probs = np.divide(raw, s, out=np.ones_like(raw) / max(M, 1), where=(s > 0))
+    return probs
+
+
+EXPLAINER = None
+if SHAP_AVAILABLE:
+    try:
+        # Background: 50 "average" institutions at z=0
+        BACKGROUND_Z = np.zeros((50, len(FEATURES)), dtype=np.float32)
+        EXPLAINER = shap.KernelExplainer(model_proba_from_z, BACKGROUND_Z)
+        print("SHAP KernelExplainer initialized.")
+    except Exception as e:
+        EXPLAINER = None
+        print("⚠️  Failed to initialize SHAP explainer:", repr(e))
+else:
+    print("SHAP not installed; explanations disabled.")
+
+
 # ----------------- FastAPI -----------------
-app = FastAPI(title="Static Fingerprint API", version="1.1.0")
+app = FastAPI(title="Static Fingerprint API", version="1.2.0")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -246,6 +287,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+
 @app.get("/")
 def root():
     return {
@@ -253,6 +295,7 @@ def root():
         "try": ["GET /health", "POST /predict", "POST /debug/z"],
     }
 
+
 @app.get("/health")
 def health():
     stats_keys = []
@@ -271,8 +314,10 @@ def health():
         "imputer": bool(imputer),
         "scaler": bool(scaler),
         "stats_available": bool(stats),
+        "shap_available": bool(EXPLAINER is not None),
     }
 
+
 @app.post("/debug/z")
 async def debug_z(req: Request):
     try:
@@ -299,11 +344,14 @@ async def debug_z(req: Request):
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e), "trace": traceback.format_exc()})
 
+
 @app.post("/predict")
 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()
@@ -320,12 +368,11 @@ async def predict(req: Request):
         raw_logits = model.predict(X, verbose=0)
         probs, mode = decode_logits(raw_logits)
 
-        # Package response
         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])]
 
-        return {
+        resp: Dict[str, Any] = {
             "input_ok": (len(missing) == 0),
             "missing": missing,
             "preprocess": {
@@ -342,5 +389,26 @@ async def predict(req: Request):
                 "raw_first_row": [float(v) for v in raw_logits[0]],
             },
         }
+
+        # ---- SHAP explanation for predicted class ----
+        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))
+                    }
+                else:
+                    resp["shap_error"] = "Unexpected SHAP output shape."
+            except Exception as e:
+                resp["shap_error"] = f"SHAP computation failed: {repr(e)}"
+        else:
+            resp["shap_error"] = "SHAP not available in this environment."
+
+        return resp
+
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e), "trace": traceback.format_exc()})