Update model training

bodywork-ml · Jun 20, 2022 · 9bfc942 · 9bfc942
1 parent bcc5e7d
commit 9bfc942
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/bodywork.yaml b/bodywork.yaml
@@ -13,8 +13,8 @@ stages:
       - pandas>=1.2.5
       - scikit-learn>=0.24.2
       - git+https://github.com/bodywork-ml/[email protected]
-    cpu_request: 0.5
-    memory_request_mb: 100
+    cpu_request: 1.0
+    memory_request_mb: 500
     batch:
       max_completion_time_seconds: 120
       retries: 2

diff --git a/pipeline/train_model.py b/pipeline/train_model.py
@@ -18,7 +18,7 @@
 PRODUCT_CODE_MAP = {"SKU001": 0, "SKU002": 1, "SKU003": 2, "SKU004": 3, "SKU005": 4}
 HYPERPARAM_GRID = {
     "random_state": [42],
-    "criterion": ["mse", "mae"],
+    "criterion": ["squared_error", "absolute_error"],
     "max_depth": [2, 3, 4, 5, 6, 7, 8, 9, 10, None],
     "min_samples_split": [2, 3, 4, 5, 6, 7, 8, 9, 10],
     "min_samples_leaf": [2, 3, 4, 5, 6, 7, 8, 9, 10],

diff --git a/requirements_pipe.txt b/requirements_pipe.txt
@@ -1,6 +1,6 @@
 numpy>=1.21.0
 pandas>=1.2.5
-scikit-learn>=0.24.2
+scikit-learn>=1.0.0
 boto3>=1.17.101
 joblib>=1.0.1
 fastapi>=0.65.2