chore(shap): sample background dataset to make explanation faster

opentargets · Feb 20, 2025 · ade23c5 · ade23c5
1 parent 2f86159
commit ade23c5
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/gentropy/dataset/l2g_prediction.py b/src/gentropy/dataset/l2g_prediction.py
@@ -246,9 +246,11 @@ def _explain(
             raise AttributeError(
                 "`model.training_data` is missing, seed dataset to get shapley values cannot be created."
             )
-        background_data = model.training_data._df.select(
-            *model.features_list
-        ).toPandas()
+        background_data = (
+            model.training_data._df.select(*model.features_list)
+            .toPandas()
+            .sample(n=1_000)
+        )
         explainer = shap.TreeExplainer(
             model.model,
             data=background_data,