feat: add optional resampler= kwarg to AutoML.fit for per-fold class-imbalance handling (#1200)

flaml/automl/automl.py

@@ -19,7 +19,13 @@
 
 from flaml import tune
 from flaml.automl.logger import logger, logger_formatter
-from flaml.automl.ml import huggingface_metric_to_mode, sklearn_metric_name_set, spark_metric_name_dict, train_estimator
+from flaml.automl.ml import (
+    _resample_training_data,
+    huggingface_metric_to_mode,
+    sklearn_metric_name_set,
+    spark_metric_name_dict,
+    train_estimator,
+)
 from flaml.automl.spark import DataFrame, Series, psDataFrame, psSeries
 from flaml.automl.state import AutoMLState, SearchState
 from flaml.automl.task.factory import task_factory
@@ -1845,6 +1851,8 @@ def fit(
         mlflow_logging=None,
         fit_kwargs_by_estimator=None,
         mlflow_exp_name=None,
+        *,
+        resampler=None,
         **fit_kwargs,
     ):
         """Find a model for a given task.
@@ -2163,6 +2171,15 @@ def cv_score_agg_func(val_loss_folds, log_metrics_folds):
         }
         ```
 
+            resampler: object, default=None | An imbalanced-learn-compatible resampler
+                (such as `imblearn.over_sampling.SMOTE`) that is cloneable via
+                `sklearn.base.clone` and exposes `fit_resample(X, y) -> (X, y)`. When set,
+                the resampler is cloned and applied to each cross-validation fold's or
+                holdout's training partition and to final, retrain, and ensemble training
+                data. Validation partitions are left at the raw class distribution. Not
+                compatible with `sample_weight` (resampling breaks the 1-to-1 row alignment
+                with weights); passing both raises `ValueError`. Off by default. See issue
+                #1200 for the design discussion and benchmarks.
             **fit_kwargs: Other key word arguments to pass to fit() function of
                 the searched learners, such as sample_weight. Below are a few examples of
                 estimator-specific parameters:
@@ -2336,6 +2353,31 @@ def cv_score_agg_func(val_loss_folds, log_metrics_folds):
             self._state.resources_per_trial = {"cpu": n_jobs} if n_jobs > 0 else {"cpu": 1}
         self._state.free_mem_ratio = self._settings.get("free_mem_ratio") if free_mem_ratio is None else free_mem_ratio
         self._state.task = task
+        fit_kwargs_by_estimator = fit_kwargs_by_estimator or self._settings.get("fit_kwargs_by_estimator")
+        if resampler is not None:
+            weight_sources = [fit_kwargs] + list((fit_kwargs_by_estimator or {}).values())
+            if any("sample_weight" in kw for kw in weight_sources):
+                raise ValueError(
+                    "Cannot combine 'resampler' with 'sample_weight' (including via "
+                    "fit_kwargs_by_estimator) — resampling breaks the 1-to-1 row alignment "
+                    "with sample weights. Use either resampling or sample weighting, not both."
+                )
+            if not callable(getattr(resampler, "fit_resample", None)):
+                raise TypeError(
+                    "'resampler' must expose a fit_resample(X, y) -> (X, y) method "
+                    "(e.g., an imbalanced-learn BaseSampler such as SMOTE)."
+                )
+            try:
+                from sklearn.base import clone
+
+                clone(resampler)
+            except Exception as e:
+                raise TypeError(
+                    "'resampler' must be cloneable via sklearn.base.clone (implement "
+                    "get_params/set_params, e.g. by subclassing sklearn.base.BaseEstimator); "
+                    f"cloning failed with: {e}"
+                ) from e
+        task._resampler = resampler
         self._state.log_training_metric = log_training_metric
 
         self._state.fit_kwargs = fit_kwargs
@@ -2348,7 +2390,6 @@ def cv_score_agg_func(val_loss_folds, log_metrics_folds):
             if mlflow_logging is None
             else mlflow_logging
         )
-        fit_kwargs_by_estimator = fit_kwargs_by_estimator or self._settings.get("fit_kwargs_by_estimator")
         self._state.fit_kwargs_by_estimator = fit_kwargs_by_estimator.copy()  # shallow copy of fit_kwargs_by_estimator
         self._state.weight_val = sample_weight_val
         self._mlflow_exp_name = mlflow_exp_name
@@ -3299,15 +3340,20 @@ def _search(self):
                 sample_weight_dict = (
                     (self._sample_weight_full is not None) and {"sample_weight": self._sample_weight_full} or {}
                 )
+                X_ensemble_train, y_ensemble_train = _resample_training_data(
+                    self._X_train_all,
+                    self._y_train_all,
+                    self._state.task,
+                )
                 for e in estimators:
                     e[1].__class__.init()
                 import joblib
 
                 try:
                     logger.info("Building ensemble with tuned estimators")
                     stacker.fit(
-                        self._X_train_all,
-                        self._y_train_all,
+                        X_ensemble_train,
+                        y_ensemble_train,
                         **sample_weight_dict,  # NOTE: _search is after kwargs is updated to fit_kwargs_by_estimator
                     )
                     logger.info(f"ensemble: {stacker}")
@@ -3325,8 +3371,8 @@ def _search(self):
                             passthrough=False,
                         )
                         stacker.fit(
-                            self._X_train_all,
-                            self._y_train_all,
+                            X_ensemble_train,
+                            y_ensemble_train,
                             **sample_weight_dict,  # NOTE: _search is after kwargs is updated to fit_kwargs_by_estimator
                         )
                         logger.info(f"ensemble: {stacker}")

flaml/automl/ml.py

@@ -448,13 +448,24 @@ def train_estimator(
         fit_kwargs["metric"] = eval_metric
 
     if X_train is not None:
+        X_train, y_train = _resample_training_data(X_train, y_train, task)
         train_time = estimator.fit(X_train, y_train, budget=budget, free_mem_ratio=free_mem_ratio, **fit_kwargs)
     else:
         estimator = estimator.estimator_class(**estimator.params)
     train_time = time.time() - start_time
     return estimator, train_time
 
 
+def _resample_training_data(X_train, y_train, task):
+    resampler = getattr(task, "_resampler", None)
+    if resampler is None:
+        return X_train, y_train
+
+    from sklearn.base import clone
+
+    return clone(resampler).fit_resample(X_train, y_train)
+
+
 def norm_confusion_matrix(y_true: Union[np.array, Series], y_pred: Union[np.array, Series]):
     """normalized confusion matrix.
 
@@ -525,6 +536,7 @@ def get_val_loss(
     #     fit_kwargs['groups_val'] = groups_val
     #     fit_kwargs['X_val'] = X_val
     #     fit_kwargs['y_val'] = y_val
+    X_train, y_train = _resample_training_data(X_train, y_train, task)
     estimator.fit(X_train, y_train, budget=budget, free_mem_ratio=free_mem_ratio, **fit_kwargs)
     val_loss, metric_for_logging, pred_time, _ = _eval_estimator(
         config,