Update tools/explain_tool.py

tools/explain_tool.py

@@ -17,8 +17,7 @@ from utils.tracing import Tracer
 
 class ExplainTool:
     """
-    Generates lightweight global SHAP visualizations (bar + beeswarm) for a sample
-    of the current DataFrame. Designed to run on CPU in HF Spaces.
+    Generates global SHAP visualizations for a sample of rows (CPU-friendly).
     """
     def __init__(self, cfg: AppConfig, tracer: Tracer):
         self.cfg = cfg
@@ -32,20 +31,11 @@ class ExplainTool:
         token = os.getenv("HF_TOKEN")
         repo = self.cfg.hf_model_repo
 
-        model_path = hf_hub_download(
-            repo_id=repo,
-            filename="model.pkl",
-            token=token
-        )
+        model_path = hf_hub_download(repo_id=repo, filename="model.pkl", token=token)
         self._model = joblib.load(model_path)
 
-        # read optional feature metadata to keep column order consistent
         try:
-            meta_path = hf_hub_download(
-                repo_id=repo,
-                filename="feature_metadata.json",
-                token=token
-            )
+            meta_path = hf_hub_download(repo_id=repo, filename="feature_metadata.json", token=token)
             with open(meta_path, "r", encoding="utf-8") as f:
                 meta = json.load(f) or {}
             self._feature_order = meta.get("feature_order")
@@ -61,38 +51,26 @@ class ExplainTool:
         return "data:image/png;base64," + base64.b64encode(buf.read()).decode()
 
     def run(self, df: Optional[pd.DataFrame]) -> Dict[str, str]:
-        """
-        Returns dict of {plot_name: data_uri_png}. If df is None/empty, returns {}.
-        """
         self._ensure_model()
         if df is None or len(df) == 0:
             return {}
 
-        # Select & sample features
         if self._feature_order:
-            missing = [c for c in self._feature_order if c not in df.columns]
-            if missing:
-                # best effort: intersect
-                X = df[[c for c in self._feature_order if c in df.columns]].copy()
-            else:
-                X = df[self._feature_order].copy()
+            cols = [c for c in self._feature_order if c in df.columns]
+            X = df[cols].copy()
         else:
             X = df.copy()
 
-        # Small sample for speed
         n = min(len(X), 500)
         sample = X.sample(n, random_state=42) if len(X) > n else X
 
-        # Build explainer and compute SHAP values
         explainer = shap.Explainer(self._model, sample)
         sv = explainer(sample)
 
-        # --- Global bar plot ---
         fig_bar = plt.figure()
         shap.plots.bar(sv, show=False)
         bar_uri = self._to_data_uri(fig_bar)
 
-        # --- Beeswarm plot ---
         fig_bee = plt.figure()
         shap.plots.beeswarm(sv, show=False)
         bee_uri = self._to_data_uri(fig_bee)