diff --git a/modules/ml-pipeline/src/pipeline/core/MLModels.py b/modules/ml-pipeline/src/pipeline/core/MLModels.py index dabe154..2444314 100644 --- a/modules/ml-pipeline/src/pipeline/core/MLModels.py +++ b/modules/ml-pipeline/src/pipeline/core/MLModels.py @@ -11,9 +11,6 @@ import joblib import pandas as pd from pathlib import Path from typing import Union, List -from sklearn import linear_model -from sklearn.svm import SVR -from autogluon.tabular import TabularDataset, TabularPredictor from core.interface.InterfaceModels import MLModel from core.Logger import logger @@ -69,6 +66,8 @@ class SKLearnLinearRegression: """ Method to train a model """ + from sklearn import linear_model + self.model = linear_model.LinearRegression() x_train = data.iloc[:, data.columns != target] @@ -117,6 +116,7 @@ class SKLearnSVMRegression: """ Method to train a model """ + from sklearn.svm import SVR validate_dict_keys( list(model_hyperparameters.keys()), @@ -161,6 +161,8 @@ class AutogluonAutoML: """ Method to load a model """ + from autogluon.tabular import TabularPredictor + filepath = str(path) self.model = TabularPredictor.load(path=filepath) @@ -186,6 +188,8 @@ class AutogluonAutoML: """ Method to train a model """ + from autogluon.tabular import TabularDataset, TabularPredictor + # Force Parallel Model fitting os.environ["AG_FORCE_PARALLEL"] = "True"