fix: 🐛 fix fasttreeshap for Leshy, try import only if fastshap is chosen

ThomasBury · Nov 28, 2023 · 0de6a86 · 0de6a86
1 parent 02616fa
commit 0de6a86
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/src/arfs/feature_selection/allrelevant.py b/src/arfs/feature_selection/allrelevant.py
@@ -314,11 +314,12 @@ def fit(self, X, y, sample_weight=None):
             Nothing but attributes
 
         """
-        try:
-            from fasttreeshap import TreeExplainer as FastTreeExplainer
-        except ImportError:
-            warnings.warn("fasttreeshap is not installed. Fallback to shap.")
-            self.importance = "shap"
+        if self.importance == "fastshap":
+            try:
+                from fasttreeshap import TreeExplainer as FastTreeExplainer
+            except ImportError:
+                warnings.warn("fasttreeshap is not installed. Fallback to shap.")
+                self.importance = "shap"
 
         if isinstance(X, pd.DataFrame):
             self.feature_names_in_ = X.columns.to_numpy()