add self.classes_

koaning · Jan 2, 2024 · 5506794 · 5506794
1 parent 6c7501a
commit 5506794
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/hulearn/__init__.py b/hulearn/__init__.py
@@ -1 +1 @@
-__version__ = "0.3.4"
+__version__ = "0.3.5"
diff --git a/hulearn/classification/functionclassifier.py b/hulearn/classification/functionclassifier.py
@@ -1,6 +1,6 @@
 from sklearn.base import BaseEstimator, ClassifierMixin
 from sklearn.utils.validation import check_is_fitted
-
+from sklearn.utils.multiclass import unique_labels
 
 class FunctionClassifier(BaseEstimator, ClassifierMixin):
     """
@@ -46,6 +46,7 @@ def fit(self, X, y):
         """
         # Run it to confirm no error happened.
         _ = self.func(X, **self.kwargs)
+        self.classes_ = unique_labels(y)
         self.fitted_ = True
         return self
 
@@ -55,6 +56,7 @@ def partial_fit(self, X, y, classes=None, sample_weight=None):
         """
         # Run it to confirm no error happened.
         _ = self.func(X, **self.kwargs)
+        self.classes_ = classes
         self.fitted_ = True
         return self
 

diff --git a/setup.py b/setup.py
@@ -68,6 +68,7 @@ def read(fname):
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
         "License :: OSI Approved :: MIT License",
         "Topic :: Scientific/Engineering",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",