Update app.py

app.py

@@ -118,19 +118,21 @@ def explainability(_):
 
     shap_path = None
     if isinstance(shap_values, list):
-    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()
+        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"
-    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,