From 42d1dcfa0e16bb35a7dd05065beb0758b57eb78e Mon Sep 17 00:00:00 2001 From: Daniel Grindrod Date: Tue, 17 Dec 2024 05:54:49 +0000 Subject: [PATCH] fix: Fixed bug with catboost and groups (#1383) Co-authored-by: Daniel Grindrod --- flaml/automl/model.py | 4 ++-- test/automl/test_split.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/flaml/automl/model.py b/flaml/automl/model.py index 024a641dd..860165fba 100644 --- a/flaml/automl/model.py +++ b/flaml/automl/model.py @@ -2066,8 +2066,8 @@ class CatBoostEstimator(BaseEstimator): self.estimator_class = CatBoostRegressor def fit(self, X_train, y_train, budget=None, free_mem_ratio=0, **kwargs): - if "is_retrain" in kwargs: - kwargs.pop("is_retrain") + kwargs.pop("is_retrain", None) + kwargs.pop("groups", None) start_time = time.time() deadline = start_time + budget if budget else np.inf train_dir = f"catboost_{str(start_time)}" diff --git a/test/automl/test_split.py b/test/automl/test_split.py index 33994a883..6c8e678b2 100644 --- a/test/automl/test_split.py +++ b/test/automl/test_split.py @@ -68,7 +68,7 @@ def test_groups(): "model_history": True, "eval_method": "cv", "groups": np.random.randint(low=0, high=10, size=len(y)), - "estimator_list": ["lgbm", "rf", "xgboost", "kneighbor"], + "estimator_list": ["catboost", "lgbm", "rf", "xgboost", "kneighbor"], "learner_selector": "roundrobin", } automl.fit(X, y, **automl_settings) @@ -108,6 +108,7 @@ def test_stratified_groupkfold(): "split_type": splitter, "groups": X_train["Airline"], "estimator_list": [ + "catboost", "lgbm", "rf", "xgboost",