diff --git a/flaml/automl/task/generic_task.py b/flaml/automl/task/generic_task.py index 0726332ba2..3c7cc7a1ed 100644 --- a/flaml/automl/task/generic_task.py +++ b/flaml/automl/task/generic_task.py @@ -964,6 +964,20 @@ def prepare_data( y_train, y_val = y_train_all[train_idx], y_train_all[val_idx] state.groups = state.groups_all[train_idx] state.groups_val = state.groups_all[val_idx] + if "sample_weight" in state.fit_kwargs: + # NOTE: _prepare_data is before kwargs is updated to fit_kwargs_by_estimator + weight = state.fit_kwargs["sample_weight"] + if weight is None: + pass + elif isinstance(weight, pd.Series): + state.fit_kwargs["sample_weight"] = weight.iloc[train_idx] + state.weight_val = weight.iloc[val_idx] + else: + # train_idx/val_idx are ndarrays from GroupShuffleSplit, so a plain + # list or tuple weight needs converting before it supports fancy indexing + weight = np.asarray(weight) + state.fit_kwargs["sample_weight"] = weight[train_idx] + state.weight_val = weight[val_idx] elif self.is_classification(): # for classification, make sure the labels are complete in both # training and validation data diff --git a/test/automl/test_split.py b/test/automl/test_split.py index 8eba4bf28d..36d79fee87 100644 --- a/test/automl/test_split.py +++ b/test/automl/test_split.py @@ -74,6 +74,112 @@ def test_time_split_with_sample_weight(): assert automl.model is not None +def test_group_split_with_sample_weight_alignment(): + """Regression for group holdout: sample_weight must be split by the group + train/val indices, not left full-length and sliced positionally downstream.""" + from unittest.mock import patch + + from sklearn.ensemble import RandomForestClassifier + + n = 200 + rng = np.random.default_rng(42) + # column 0 is a row id, so each row's expected weight is recoverable from X + X = np.column_stack([np.arange(n, dtype=float), rng.normal(size=n)]) + groups = np.repeat(np.arange(20), 10) + y = (groups % 2).astype(int) + sample_weight = 1000.0 + np.arange(n, dtype=float) + + captured = [] + original_fit = RandomForestClassifier.fit + + def spy_fit(self, X, y, sample_weight=None, **kwargs): + captured.append( + ( + np.asarray(X).copy(), + None if sample_weight is None else np.asarray(sample_weight, dtype=float).copy(), + ) + ) + return original_fit(self, X, y, sample_weight=sample_weight, **kwargs) + + automl = AutoML() + with patch.object(RandomForestClassifier, "fit", spy_fit): + automl.fit( + X_train=X, + y_train=y, + sample_weight=sample_weight, + groups=groups, + split_type="group", + eval_method="holdout", + task="classification", + time_budget=-1, + max_iter=2, + estimator_list=["rf"], + ) + assert automl.model is not None + assert captured + for X_fit, weight_fit in captured: + assert weight_fit is not None + # each row trains with its own weight (1000 + row id) after the split + np.testing.assert_array_equal(weight_fit, 1000.0 + np.asarray(X_fit)[:, 0]) + # the validation loss must be weighted as well + assert automl._state.weight_val is not None + + +def test_group_split_with_sample_weight_series_alignment(): + """Same regression as above, but sample_weight is a pd.Series with a + non-default index, to verify the group branch slices it positionally + (.iloc) rather than by label.""" + from unittest.mock import patch + + from sklearn.ensemble import RandomForestClassifier + + n = 200 + rng = np.random.default_rng(7) + X = pd.DataFrame( + {"id": np.arange(n, dtype=float), "feat": rng.normal(size=n)}, + ) + groups = np.repeat(np.arange(20), 10) + y = pd.Series((groups % 2).astype(int)) + # non-default index: label != positional order, so .loc and .iloc would disagree + weight_index = np.arange(n) + 500 + sample_weight = pd.Series(1000.0 + np.arange(n, dtype=float), index=weight_index) + + captured = [] + original_fit = RandomForestClassifier.fit + + def spy_fit(self, X, y, sample_weight=None, **kwargs): + captured.append( + ( + np.asarray(X).copy(), + None if sample_weight is None else np.asarray(sample_weight, dtype=float).copy(), + ) + ) + return original_fit(self, X, y, sample_weight=sample_weight, **kwargs) + + automl = AutoML() + with patch.object(RandomForestClassifier, "fit", spy_fit): + automl.fit( + X_train=X, + y_train=y, + sample_weight=sample_weight, + groups=groups, + split_type="group", + eval_method="holdout", + task="classification", + time_budget=-1, + max_iter=2, + estimator_list=["rf"], + ) + assert automl.model is not None + assert captured + for X_fit, weight_fit in captured: + assert weight_fit is not None + # each row trains with its own weight (1000 + row id) after the split, + # regardless of the Series' original (non-default) index + np.testing.assert_array_equal(weight_fit, 1000.0 + np.asarray(X_fit)[:, 0]) + assert automl._state.weight_val is not None + + def test_groups_for_classification_task(): from sklearn.externals._arff import ArffException