Merge pull request #36 from Tihsrah/Tihsrah

Adding clustering 5 models inside single_model_dict #32
sagnik1511 · Feb 12, 2022 · 2d790e3 · 2d790e3
2 parents 6e77dc7 + fff70d2
commit 2d790e3
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/tab_automl/automl/models.py b/tab_automl/automl/models.py
@@ -9,7 +9,11 @@
 from sklearn.tree import DecisionTreeClassifier
 from sklearn.neighbors import KNeighborsClassifier
 from sklearn.neighbors import KNeighborsRegressor
-
+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": {
@@ -26,5 +30,12 @@
         "Random Forest Classifier": RandomForestClassifier,
         "XGBoost Classifier": XGBClassifier,
         "KNN Classifier": KNeighborsClassifier
+    },
+    "clustering": {
+        "Affinity Propagation": AffinityPropagation,
+        "Agglomerative Clustering": AgglomerativeClustering,
+        "Birch": Birch,
+        "DBSCAN": DBSCAN,
+        "KMeans": KMeans
     }
 }