From 84e66d3b1376aef810dfbb53d4f7da7a2037673a Mon Sep 17 00:00:00 2001 From: Tihsrah Date: Fri, 11 Feb 2022 15:13:54 +0530 Subject: [PATCH] [This PR adds clustering models {#32}] Add clustering models inside single_model_dict --- tab_automl/automl/models.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/tab_automl/automl/models.py b/tab_automl/automl/models.py index a35b03d..d2181d0 100644 --- a/tab_automl/automl/models.py +++ b/tab_automl/automl/models.py @@ -6,17 +6,30 @@ from sklearn.linear_model import Lasso, LinearRegression, Ridge from sklearn.tree import DecisionTreeClassifier +from sklearn.cluster import AffinityPropagation +from sklearn.cluster import AgglomerativeClustering +from sklearn.cluster import Birch +from sklearn.cluster import DBSCAN +from sklearn.cluster import KMeans single_model_dict = { "regression": { "Linear Regression": LinearRegression, "Lasso Regression": Lasso, "Ridge Regression": Ridge, - "Random Forest Regression": RandomForestRegressor + "Random Forest Regression": RandomForestRegressor, }, "classification": { "Decision Tree Classifier": DecisionTreeClassifier, "Light Gradient Boosting Classifier": LGBMClassifier, "Random Forest Classifier": RandomForestClassifier + }, + "clustering":{ + "AffinityPropagation":AffinityPropagation, + "AgglomerativeClustering":AgglomerativeClustering, + "Birch":Birch, + "DBSCAN":DBSCAN, + "KMeans":KMeans + } }