From ea325745aa2e156babbb2db045719845aea15794 Mon Sep 17 00:00:00 2001 From: PalaVenkiReddy Date: Thu, 16 Feb 2023 19:29:47 +0530 Subject: [PATCH] Added new models for Regression Training --- tab_automl/automl/models.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tab_automl/automl/models.py b/tab_automl/automl/models.py index 30fa2b9..fc65bd2 100644 --- a/tab_automl/automl/models.py +++ b/tab_automl/automl/models.py @@ -1,7 +1,7 @@ """ This files holds all codes for defined and used models through the life cycle of the package. """ -from lightgbm import LGBMClassifier +from lightgbm import LGBMClassifier, LGBMRegressor from xgboost import XGBClassifier from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor from sklearn.linear_model import Lasso, LinearRegression, Ridge @@ -19,6 +19,7 @@ from sklearn.cluster import OPTICS from sklearn.cluster import SpectralClustering from sklearn.mixture import GaussianMixture +from sklearn.tree import DecisionTreeRegressor single_model_dict = { "regression": { @@ -27,7 +28,9 @@ "Ridge Regression": Ridge, "Random Forest Regression": RandomForestRegressor, "Support Vector Regression": SVR, - "KNN Regressor": KNeighborsRegressor + "KNN Regressor": KNeighborsRegressor, + "Decision Tree Regressor": DecisionTreeRegressor, + "Light Gradient Boosting Regressor": LGBMRegressor, }, "classification": { "Decision Tree Classifier": DecisionTreeClassifier,