Merge pull request #36 from TobyBoyne/master

Fix categorical feature bug with LGBM v4.1+
cog-imperial · Jun 21, 2024 · 5d7003d · 5d7003d
2 parents 7e2e50e + ead026a
commit 5d7003d
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 4 deletions.
diff --git a/entmoot/models/mean_models/tree_ensemble.py b/entmoot/models/mean_models/tree_ensemble.py
@@ -84,15 +84,14 @@ def _train_lgbm(self, X, y):
                 train_data = lgb.Dataset(
                     X,
                     label=y,
-                    categorical_feature=self._problem_config.cat_idx,
+                    categorical_feature=list(self._problem_config.cat_idx),
                     free_raw_data=False,
                     params={"verbose": -1},
                 )
 
                 tree_model = lgb.train(
                     self._train_params,
                     train_data,
-                    categorical_feature=self._problem_config.cat_idx,
                     #verbose_eval=False,
                 )
             else:

diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -1,5 +1,5 @@
 numpy<2.0.0
-lightgbm==4.0.0
+lightgbm>=4.0.0
 pyomo
 gurobipy
 pytest-cov

diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
 numpy<2.0.0
-lightgbm==4.0.0
+lightgbm>=4.0.0
 gurobipy
 pyomo