Update app.py

app.py

@@ -118,21 +118,19 @@ def explainability(_):
 
     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
+    class_idx = 0  # or dynamically pick class with most samples
+    shap.summary_plot(shap_values[class_idx], X_test, show=False)
+    shap_path = f"./shap_class_{class_idx}.png"
+    plt.title(f"SHAP Summary - Class {class_idx}")
+    plt.savefig(shap_path)
+    wandb.log({f"shap_class_{class_idx}": wandb.Image(shap_path)})
+    plt.clf()
     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()
+    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()
 
     lime_explainer = lime.lime_tabular.LimeTabularExplainer(
         X_train.values,