Update app.py

app.py

@@ -119,36 +119,34 @@ def explainability(_):
     shap_values = explainer.shap_values(X_test)
 
     try:
-        if isinstance(shap_values, list):  # Multiclass
+        if isinstance(shap_values, list):
             class_idx = 0
-            X_shap = pd.DataFrame(
-                X_test.values[:, :shap_values[class_idx].shape[1]],
-                columns=X_test.columns[:shap_values[class_idx].shape[1]]
-            )
-            shap.summary_plot(shap_values[class_idx], X_shap, show=False)
-            shap_path = f"./shap_class_{class_idx}.png"
-            plt.title(f"SHAP Summary - Class {class_idx}")
+            sv = shap_values[class_idx]
         else:
-            X_shap = pd.DataFrame(
-                X_test.values[:, :shap_values.shape[1]],
-                columns=X_test.columns[:shap_values.shape[1]]
-            )
-            shap.summary_plot(shap_values, X_shap, show=False)
-            shap_path = "./shap_plot.png"
+            sv = shap_values
 
+        # Align number of columns
+        if sv.shape[1] != X_test.shape[1]:
+            X_test_trimmed = pd.DataFrame(X_test.values[:, :sv.shape[1]], columns=X_test.columns[:sv.shape[1]])
+        else:
+            X_test_trimmed = X_test
+
+        shap.summary_plot(sv, X_test_trimmed, show=False)
+        shap_path = "./shap_plot.png"
+        plt.title("SHAP Summary")
         plt.savefig(shap_path)
-        if wandb.run is not None:
+        if wandb.run:
             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}")
+        print("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.text(0.5, 0.5, f"SHAP Error:\n{str(e)}", ha='center', va='center')
         plt.axis('off')
         plt.savefig(shap_path)
-        if wandb.run is not None:
+        if wandb.run:
             wandb.log({"shap_error": wandb.Image(shap_path)})
         plt.clf()
 
@@ -163,7 +161,7 @@ def explainability(_):
     lime_fig = lime_exp.as_pyplot_figure()
     lime_path = "./lime_plot.png"
     lime_fig.savefig(lime_path)
-    if wandb.run is not None:
+    if wandb.run:
         wandb.log({"lime_explanation": wandb.Image(lime_path)})
     plt.clf()
 
@@ -171,6 +169,7 @@ def explainability(_):
 
 
 
+
 with gr.Blocks() as demo:
     gr.Markdown("## 📊 AI-Powered Data Analysis with Hyperparameter Optimization")