Merge pull request #37 from snega16/snega16

[Issue {#32}]Add clustering models inside single_model_dict
sagnik1511 · Feb 13, 2022 · 9f6259c · 9f6259c
2 parents 2d790e3 + 60f3da6
commit 9f6259c
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/tab_automl/automl/models.py b/tab_automl/automl/models.py
@@ -14,6 +14,11 @@
 from sklearn.cluster import Birch
 from sklearn.cluster import DBSCAN
 from sklearn.cluster import KMeans
+from sklearn.cluster import MiniBatchKMeans
+from sklearn.cluster import MeanShift
+from sklearn.cluster import OPTICS
+from sklearn.cluster import SpectralClustering
+from sklearn.mixture import GaussianMixture
 
 single_model_dict = {
     "regression": {
@@ -36,6 +41,11 @@
         "Agglomerative Clustering": AgglomerativeClustering,
         "Birch": Birch,
         "DBSCAN": DBSCAN,
-        "KMeans": KMeans
+        "KMeans": KMeans,
+        "MiniBatchKMeans": MiniBatchKMeans,
+        "MeanShift": MeanShift,
+        "OPTICS": OPTICS,
+        "SpectralClustering": SpectralClustering,
+        "GaussianMixture": GaussianMixture
     }
 }