Adding parameter for regression metric and time limit.

georgia-tech-db · Oct 18, 2023 · 37c94f8 · 37c94f8
1 parent 8e1bd05
commit 37c94f8
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/evadb/configuration/constants.py b/evadb/configuration/constants.py
@@ -34,3 +34,4 @@
 DEFAULT_TRAIN_TIME_LIMIT = 120
 DEFAULT_DOCUMENT_CHUNK_SIZE = 4000
 DEFAULT_DOCUMENT_CHUNK_OVERLAP = 200
+DEFAULT_TRAIN_REGRESSION_METRIC = "rmse"
diff --git a/evadb/executor/create_function_executor.py b/evadb/executor/create_function_executor.py
@@ -25,6 +25,7 @@
 from evadb.catalog.models.function_io_catalog import FunctionIOCatalogEntry
 from evadb.catalog.models.function_metadata_catalog import FunctionMetadataCatalogEntry
 from evadb.configuration.constants import (
+    DEFAULT_TRAIN_REGRESSION_METRIC,
     DEFAULT_TRAIN_TIME_LIMIT,
     EvaDB_INSTALLATION_DIR,
 )
@@ -190,7 +191,7 @@ def handle_xgboost_function(self):
         model = AutoML()
         settings = {
             "time_budget": arg_map.get("time_limit", DEFAULT_TRAIN_TIME_LIMIT),
-            "metric": "r2",
+            "metric": arg_map.get("metric", DEFAULT_TRAIN_REGRESSION_METRIC),
             "estimator_list": ["xgboost"],
             "task": "regression",
         }

diff --git a/test/integration_tests/long/test_model_train.py b/test/integration_tests/long/test_model_train.py
@@ -102,7 +102,8 @@ def test_xgboost_regression(self):
             ( SELECT number_of_rooms, number_of_bathrooms, days_on_market, rental_price FROM HomeRentals )
             TYPE XGBoost
             PREDICT 'rental_price'
-            TIME_LIMIT 180;
+            TIME_LIMIT 180
+            METRIC 'r2';
         """
         execute_query_fetch_all(self.evadb, create_predict_function)