Skip to content

Commit

Permalink
[This PR adds clustering models {sagnik1511#32}]
Browse files Browse the repository at this point in the history
Add clustering models inside single_model_dict
  • Loading branch information
Tihsrah committed Feb 11, 2022
1 parent aa31797 commit 84e66d3
Showing 1 changed file with 14 additions and 1 deletion.
15 changes: 14 additions & 1 deletion tab_automl/automl/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

}
}

0 comments on commit 84e66d3

Please sign in to comment.