Update app.py

app.py

@@ -628,6 +628,12 @@ def train_and_save(
             "selection_method": "SelectFromModel(L1 saga, threshold=median)" if use_feature_selection else None,
             "note": "If SVD is enabled, SHAP becomes component-level (less interpretable)."
         },
+        "shap_background": {
+            "file": "background.csv",
+            "max_rows": 200,
+            "note": "Raw (pre-transform) background sample for SHAP LinearExplainer."
+        },
+        
         "positive_class": str(pos_class),
         "metrics": metrics,
     }
@@ -668,6 +674,131 @@ def ensure_model_repo_exists(model_repo_id: str, token: str):
     except Exception:
         pass
 
+def coerce_X_like_schema(X: pd.DataFrame, feature_cols: list[str], num_cols: list[str], cat_cols: list[str]) -> pd.DataFrame:
+    """
+    Ensure X has correct columns and coercions, matching your training/inference convention.
+    """
+    X = X[feature_cols].copy().replace({pd.NA: np.nan})
+
+    for c in num_cols:
+        if c in X.columns:
+            X[c] = pd.to_numeric(X[c], errors="coerce")
+
+    for c in cat_cols:
+        if c in X.columns:
+            X[c] = X[c].astype("object")
+            X.loc[X[c].isna(), c] = np.nan
+            X[c] = X[c].map(lambda v: v if pd.isna(v) else str(v))
+
+    return X
+
+
+def get_shap_background_auto(model_repo_id: str, feature_cols: list[str], num_cols: list[str], cat_cols: list[str]) -> pd.DataFrame | None:
+    """
+    Attempts to load SHAP background from HF repo. Returns coerced background or None.
+    """
+    df_bg = load_latest_background(model_repo_id)
+    if df_bg is None:
+        return None
+
+    # Ensure required columns exist
+    missing = [c for c in feature_cols if c not in df_bg.columns]
+    if missing:
+        return None
+
+    return coerce_X_like_schema(df_bg, feature_cols, num_cols, cat_cols)
+
+
+
+# ============================================================
+# SHAP background persistence (best practice)
+# ============================================================
+
+def save_background_sample_csv(X_bg: pd.DataFrame, feature_cols: list[str], max_rows: int = 200, out_path: str = "background.csv"):
+    """
+    Saves a small *raw* background dataset (pre-transform) for SHAP explainer.
+    Must contain columns exactly matching feature_cols.
+    """
+    if X_bg is None or len(X_bg) == 0:
+        raise ValueError("X_bg is empty; cannot save background sample.")
+
+    X_bg = X_bg[feature_cols].copy()
+
+    if len(X_bg) > int(max_rows):
+        X_bg = X_bg.sample(int(max_rows), random_state=42)
+
+    # Preserve exact columns for future loading
+    X_bg.to_csv(out_path, index=False, encoding="utf-8")
+    return out_path
+
+
+def publish_background_to_hub(model_repo_id: str, version_tag: str, background_path: str = "background.csv"):
+    """
+    Uploads background.csv to both versioned and latest paths.
+    Requires HF_TOKEN with write permissions.
+    """
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        raise RuntimeError("HF_TOKEN not found. Add it in Space Settings → Secrets.")
+    api = HfApi(token=token)
+
+    version_bg_path = f"releases/{version_tag}/background.csv"
+
+    # Versioned
+    api.upload_file(
+        path_or_fileobj=background_path,
+        path_in_repo=version_bg_path,
+        repo_id=model_repo_id,
+        repo_type="model",
+        commit_message=f"Upload SHAP background ({version_tag})"
+    )
+
+    # Latest
+    api.upload_file(
+        path_or_fileobj=background_path,
+        path_in_repo="latest/background.csv",
+        repo_id=model_repo_id,
+        repo_type="model",
+        commit_message=f"Update latest SHAP background ({version_tag})"
+    )
+
+    return {
+        "version_bg_path": version_bg_path,
+        "latest_bg_path": "latest/background.csv",
+    }
+
+
+def load_latest_background(model_repo_id: str) -> pd.DataFrame | None:
+    """
+    Loads latest/background.csv if present. Returns None if not found / cannot load.
+    """
+    try:
+        bg_file = hf_hub_download(
+            repo_id=model_repo_id,
+            repo_type="model",
+            filename="latest/background.csv",
+        )
+        df_bg = pd.read_csv(bg_file)
+        return df_bg
+    except Exception:
+        return None
+
+
+def load_background_by_version(model_repo_id: str, version_tag: str) -> pd.DataFrame | None:
+    """
+    Loads releases/<version>/background.csv if present.
+    """
+    try:
+        bg_file = hf_hub_download(
+            repo_id=model_repo_id,
+            repo_type="model",
+            filename=f"releases/{version_tag}/background.csv",
+        )
+        df_bg = pd.read_csv(bg_file)
+        return df_bg
+    except Exception:
+        return None
+
 
 def publish_to_hub(model_repo_id: str, version_tag: str):
     """
@@ -1276,7 +1407,18 @@ with tab_train:
                     use_dimred=use_dimred,
                     svd_components=svd_components
                     )
-    
+                # --- Save background sample for SHAP (raw X_train) ---
+                try:
+                    save_background_sample_csv(
+                        X_bg=X_train,
+                        feature_cols=feature_cols,
+                        max_rows=200,
+                        out_path="background.csv"
+                    )
+                    st.success("Saved SHAP background sample (background.csv).")
+                except Exception as e:
+                    st.warning(f"Could not save SHAP background sample: {e}")
+
                 explainer = build_shap_explainer(pipe, X_train)
 
                 st.session_state.pipe = pipe
@@ -1509,6 +1651,12 @@ with tab_train:
                     try:
                         with st.spinner("Uploading to Hugging Face Model repo..."):
                             paths = publish_to_hub(MODEL_REPO_ID, version_tag)
+                            # Upload background.csv if it exists
+                            if os.path.exists("background.csv"):
+                                bg_paths = publish_background_to_hub(MODEL_REPO_ID, version_tag, background_path="background.csv")
+                                paths.update(bg_paths)
+                            else:
+                                st.warning("background.csv not found; SHAP background will not be uploaded.")
             
                         st.success("Uploaded successfully to your model repository.")
                         st.json(paths)
@@ -1576,6 +1724,21 @@ with tab_predict:
     num_cols = meta["schema"]["numeric"]
     cat_cols = meta["schema"]["categorical"]
 
+    # ------------------------------------------------------------
+    # SHAP background: prefer inference file, else HF background.csv
+    # ------------------------------------------------------------
+    df_inf = st.session_state.get("df_inf")
+    
+    if df_inf is not None:
+        # use user cohort as background (optional)
+        X_bg = coerce_X_like_schema(df_inf, feature_cols, num_cols, cat_cols)
+    else:
+        # fall back to published background
+        X_bg = get_shap_background_auto(MODEL_REPO_ID, feature_cols, num_cols, cat_cols)
+    
+    st.session_state.X_bg_for_shap = X_bg
+
+
     # 2) Now we can build lookup
     FEATURE_LOOKUP = {norm_col(c): c for c in feature_cols}
 
@@ -2429,9 +2592,25 @@ with tab_predict:
                 X_batch_t = transform_before_clf(pipe, X_batch)
                 
                 explainer = st.session_state.get("explainer")
-                if explainer is None:
-                    st.session_state.explainer = build_shap_explainer(pipe, X_inf)
+                explainer_sig = st.session_state.get("explainer_sig")
+                
+                # Create a simple signature that changes if model changes or background changes
+                # (using version + number of background rows is usually enough)
+                current_sig = (
+                    selected,  # or meta.get("created_at_utc") or meta.get("metrics", {}).get("roc_auc")
+                    None if st.session_state.get("X_bg_for_shap") is None else int(len(st.session_state["X_bg_for_shap"]))
+                )
+                
+                if explainer is None or explainer_sig != current_sig:
+                    X_bg = st.session_state.get("X_bg_for_shap")
+                    if X_bg is None:
+                        st.error("SHAP background not available. Admin must publish latest/background.csv.")
+                        st.stop()
+                
+                    st.session_state.explainer = build_shap_explainer(pipe, X_bg)
+                    st.session_state.explainer_sig = current_sig
                     explainer = st.session_state.explainer
+
                 
                 shap_vals_batch = explainer.shap_values(X_batch_t)
                 if isinstance(shap_vals_batch, list):
@@ -2627,9 +2806,25 @@ with tab_predict:
             X_one_t = transform_before_clf(pipe, X_one)
         
             explainer = st.session_state.get("explainer")
-            if explainer is None:
-                st.session_state.explainer = build_shap_explainer(pipe, X_inf)
+            explainer_sig = st.session_state.get("explainer_sig")
+            
+            # Create a simple signature that changes if model changes or background changes
+            # (using version + number of background rows is usually enough)
+            current_sig = (
+                selected,  # or meta.get("created_at_utc") or meta.get("metrics", {}).get("roc_auc")
+                None if st.session_state.get("X_bg_for_shap") is None else int(len(st.session_state["X_bg_for_shap"]))
+            )
+            
+            if explainer is None or explainer_sig != current_sig:
+                X_bg = st.session_state.get("X_bg_for_shap")
+                if X_bg is None:
+                    st.error("SHAP background not available. Admin must publish latest/background.csv.")
+                    st.stop()
+            
+                st.session_state.explainer = build_shap_explainer(pipe, X_bg)
+                st.session_state.explainer_sig = current_sig
                 explainer = st.session_state.explainer
+
         
             shap_vals = explainer.shap_values(X_one_t)
             if isinstance(shap_vals, list):