feat: use dtype dic for consistency

ThomasBury · Jun 16, 2023 · 526f04b · 526f04b
1 parent 5f356dc
commit 526f04b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/arfs/gbm.py b/src/arfs/gbm.py
@@ -25,6 +25,7 @@
 import gc
 import os
 from pathlib import Path
+from arfs.utils import create_dtype_dict
 
 QUAL_COLORS = [
     (0.188235, 0.635294, 0.854902),
@@ -180,9 +181,9 @@ def fit(
             raise KeyError("Provide the objective in the params dict")
 
         if self.cat_feat == "auto":
-            self.cat_feat = list(
-                set(list(X.columns)) - set(list(X.select_dtypes(include=[np.number])))
-            )
+            dtypes_dic = create_dtype_dict(df=X, dic_keys="dtypes")
+            category_cols = dtypes_dic["cat"] + dtypes_dic["time"] + dtypes_dic["unk"]
+            self.cat_feat = category_cols if category_cols else None
 
         if not isinstance(X, (pd.Series, pd.DataFrame)):
             X = pd.DataFrame(X)