From bef5ddafb4e119ad190d87962ac4cd279e2f64e4 Mon Sep 17 00:00:00 2001
From: Animaholic <shenguohao77@gmail.com>
Date: Mon, 27 Mar 2023 13:55:52 -0400
Subject: [PATCH] retrain model

---
 flaml/automl/automl.py             |  9 +++++++++
 test/automl/test_classification.py | 32 ++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/flaml/automl/automl.py b/flaml/automl/automl.py
index 2d8ea04a1a..40d23659d1 100644
--- a/flaml/automl/automl.py
+++ b/flaml/automl/automl.py
@@ -2876,3 +2876,12 @@ def _select_estimator(self, estimator_list):
                 q += inv[i] / s
                 if p < q:
                     return estimator_list[i]
+
+    def retrain(self, X_train, y_train, config=None):
+        if config is None:
+            config = self.best_config
+
+        automl = AutoML()
+        automl.fit(X_train, y_train, **config)
+
+        return automl
diff --git a/test/automl/test_classification.py b/test/automl/test_classification.py
index adb3081662..cae5801244 100644
--- a/test/automl/test_classification.py
+++ b/test/automl/test_classification.py
@@ -408,6 +408,38 @@ def test_sparse_matrix_lr(self):
         print(automl_experiment.best_iteration)
         print(automl_experiment.best_estimator)
 
+    def test_retrain(self):
+        from flaml.automl.data import load_openml_dataset
+
+        X_train, X_test, y_train, y_test = load_openml_dataset(
+            dataset_id=187, data_dir="./"
+        )
+
+        automl_settings = {
+            "task": "classification",
+            "estimator_list": ["xgboost"],
+            # "max_iter": 1,
+        }
+
+        automl_settings["starting_points"] = {
+            "xgboost": {
+                "n_estimators": 4,
+                "max_leaves": 4,
+                "min_child_weight": 0.26208115308159446,
+                "learning_rate": 0.25912534572860507,
+                "subsample": 0.9266743941610592,
+                "colsample_bylevel": 1.0,
+                "colsample_bytree": 1.0,
+                "reg_alpha": 0.0013933617380144255,
+                "reg_lambda": 0.18096917948292954,
+                "FLAML_sample_size": 20000,
+            },
+        }
+
+        automl = AutoML()
+        automl2 = automl.retrain(X_train, y_train, automl_settings)
+        print(automl2.model.model)
+
 
 if __name__ == "__main__":
     test = TestClassification()