Merge branch 'main' into flaml-fix-groups-issue-for-regression-tasks

Author: thinkall

flaml/automl/model.py

@@ -2066,8 +2066,8 @@ def __init__(
             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)}"

test/automl/test_split.py

@@ -69,7 +69,7 @@ def test_groups_for_classification_task():
         "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)
@@ -138,6 +138,7 @@ def test_stratified_groupkfold():
         "split_type": splitter,
         "groups": X_train["Airline"],
         "estimator_list": [
+            "catboost",
             "lgbm",
             "rf",
             "xgboost",