Add SGDClassifier (#17)

* Add SGDClassifier * Fix complaining linter
pureskillgg · Jan 1, 2023 · bc71e73 · bc71e73
1 parent 4505d10
commit bc71e73
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pureskillgg_dsdk/ds_models/s3_scikit.py b/pureskillgg_dsdk/ds_models/s3_scikit.py
@@ -34,6 +34,9 @@ def _use_model(self, dataframe):
         if self._model_type == "MiniBatchKMeans":
             labels = self._loaded_model.predict(dataframe)
             return labels
+        if self._model_type == "SGDClassifier":
+            labels = self._loaded_model.predict_proba(dataframe)
+            return labels
         raise Exception(f"Unknown model_type {self._model_type}")
 
     def invoke(self, dataframe):