Merge pull request #43 from PalaVenkiReddy/PalaVenkiReddy

#7 Added new models for Classification Training
sagnik1511 · Feb 19, 2023 · cada595 · cada595
2 parents da20afc + 68af73f
commit cada595
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/tab_automl/automl/models.py b/tab_automl/automl/models.py
@@ -20,6 +20,9 @@
 from sklearn.cluster import SpectralClustering
 from sklearn.mixture import GaussianMixture
 from sklearn.tree import DecisionTreeRegressor
+from sklearn.naive_bayes import GaussianNB
+from sklearn.svm import LinearSVC
+from sklearn.svm import SVC
 
 single_model_dict = {
     "regression": {
@@ -37,7 +40,10 @@
         "Light Gradient Boosting Classifier": LGBMClassifier,
         "Random Forest Classifier": RandomForestClassifier,
         "XGBoost Classifier": XGBClassifier,
-        "KNN Classifier": KNeighborsClassifier
+        "KNN Classifier": KNeighborsClassifier,
+        "Gaussian Naive Bayes": GaussianNB,
+        "Linear Support Vector Classifier": LinearSVC,
+        "Support Vector Classifier": SVC,
     },
     "clustering": {
         "Affinity Propagation": AffinityPropagation,