Update app.py

app.py

@@ -100,6 +100,9 @@ def train_model(_):
     return metrics, top_trials
 
 def explainability(_):
+    import warnings
+    warnings.filterwarnings("ignore")
+    
     target = df_global.columns[-1]
     X = df_global.drop(target, axis=1)
     y = df_global[target]
@@ -112,41 +115,50 @@ def explainability(_):
     model = RandomForestClassifier()
     model.fit(X_train, y_train)
 
+    # SHAP Explainability
     explainer = shap.TreeExplainer(model)
     shap_values = explainer.shap_values(X_test)
 
-    shap_path = None
-    if isinstance(shap_values, list):
-        for i, class_vals in enumerate(shap_values):
-            shap.summary_plot(class_vals, X_test, show=False)
-            class_path = f"./shap_class_{i}.png"
-            plt.title(f"SHAP Summary - Class {i}")
-            plt.savefig(class_path)
-            wandb.log({f"shap_class_{i}": wandb.Image(class_path)})
-            plt.clf()
-            if shap_path is None:
-                shap_path = class_path
-    else:
-        shap.summary_plot(shap_values, X_test, show=False)
-        shap_path = "./shap_plot.png"
+    try:
+        if isinstance(shap_values, list):  # Multiclass
+            class_idx = 0
+            shap_matrix = shap_values[class_idx]  # shape: (samples, features)
+            assert shap_matrix.shape[1] == X_test.shape[1], \
+                f"SHAP mismatch: shap {shap_matrix.shape} vs X {X_test.shape}"
+            shap.summary_plot(shap_matrix, X_test, show=False)
+            shap_path = f"./shap_class_{class_idx}.png"
+            plt.title(f"SHAP Summary - Class {class_idx}")
+        else:
+            shap.summary_plot(shap_values, X_test, show=False)
+            shap_path = "./shap_plot.png"
         plt.savefig(shap_path)
         wandb.log({"shap_summary": wandb.Image(shap_path)})
         plt.clf()
+    except Exception as e:
+        shap_path = "./shap_error.png"
+        print(f"SHAP plotting failed: {e}")
+        plt.figure(figsize=(6, 3))
+        plt.text(0.5, 0.5, f"SHAP error:\n{str(e)}", ha='center', va='center')
+        plt.savefig(shap_path)
+        wandb.log({"shap_error": wandb.Image(shap_path)})
+        plt.clf()
 
+    # LIME Explainability
     lime_explainer = lime.lime_tabular.LimeTabularExplainer(
         X_train.values,
         feature_names=X_train.columns.tolist(),
-        class_names=[str(label) for label in np.unique(y_train)],
+        class_names=[str(c) for c in np.unique(y_train)],
         mode='classification'
     )
     lime_exp = lime_explainer.explain_instance(X_test.iloc[0].values, model.predict_proba)
     lime_fig = lime_exp.as_pyplot_figure()
-    lime_fig_path = "./lime_plot.png"
-    lime_fig.savefig(lime_fig_path)
-    wandb.log({"lime_explanation": wandb.Image(lime_fig_path)})
+    lime_path = "./lime_plot.png"
+    lime_fig.savefig(lime_path)
+    wandb.log({"lime_explanation": wandb.Image(lime_path)})
     plt.clf()
 
-    return shap_path, lime_fig_path
+    return shap_path, lime_path
+
 
 with gr.Blocks() as demo:
     gr.Markdown("## 📊 AI-Powered Data Analysis with Hyperparameter Optimization")