Sat, 16 Oct 2021 21:29:37 GMT

train_model.py

@@ -1,4 +1,5 @@
 # %%
+from typing import Union
 import pandas as pd
 
 # %%
@@ -50,14 +51,16 @@ train_test_data = dict(zip(data_keys, data_vals))
 
 # %%
 from sklearn.utils.validation import check_X_y
-import pickle
+import joblib
 
 from catboost import CatBoostRegressor
 from lightgbm import LGBMRegressor
 from xgboost import XGBRegressor
 
+Regressor = Union[CatBoostRegressor, LGBMRegressor, XGBRegressor]
 
-def train_model(regType):
+
+def train_model(regType: Regressor):
     X, ys = train_test_data['X_train'], train_test_data['y_train']
     check_X_y(X, ys, multi_output=True)
     models = {}
@@ -67,9 +70,7 @@ def train_model(regType):
         reg.fit(X, y, verbose=False)
         models[target_col] = reg
         print(target_col)
-    CatBoostRegressor.__name__
-    with open(f"models/model_{regType.__name__}.pickle", "wb") as f:
-        pickle.dump(models, f)
+    joblib.dump(models, f"models/{regType.__name__}.model")
 
 
 for reg in [CatBoostRegressor, LGBMRegressor, XGBRegressor]:
@@ -79,9 +80,8 @@ for reg in [CatBoostRegressor, LGBMRegressor, XGBRegressor]:
 from sklearn import metrics
 
 
-def eval_model(regType):
-    with open(f"models/model_{regType.__name__}.pickle", "rb") as f:
-        models = pickle.load(f)
+def eval_model(regType: Regressor):
+    models = joblib.load(f"models/{regType.__name__}.model")
     X, ys = train_test_data['X_test'], train_test_data['y_test']
     evals = []
     for target_col, reg in models.items():
@@ -96,5 +96,4 @@ def eval_model(regType):
 
 
 for reg in [CatBoostRegressor, LGBMRegressor, XGBRegressor]:
-    eval_model(reg)
-# %%
+    eval_model(reg)