fix: support sklearn BaseEstimators

Signed-off-by: Kushal Batra <[email protected]>
numaproj · Sep 4, 2024 · cd88862 · cd88862
1 parent eb9ceb1
commit cd88862
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/numalogic/base.py b/numalogic/base.py
@@ -16,10 +16,10 @@
 
 import numpy.typing as npt
 import pytorch_lightning as pl
-from sklearn.base import TransformerMixin, OutlierMixin
+from sklearn.base import TransformerMixin, OutlierMixin, BaseEstimator
 
 
-class BaseTransformer(TransformerMixin):
+class BaseTransformer(BaseEstimator, TransformerMixin):
     """Base class for all transformer classes."""
 
     pass
@@ -47,7 +47,7 @@ class TorchModel(pl.LightningModule, metaclass=ABCMeta):
     pass
 
 
-class BaseThresholdModel(OutlierMixin):
+class BaseThresholdModel(BaseEstimator, OutlierMixin):
     """Base class for all threshold models."""
 
     pass
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "numalogic"
-version = "0.12.4"
+version = "0.13.0"
 description = "Collection of operational Machine Learning models and tools."
 authors = ["Numalogic Developers"]
 packages = [{ include = "numalogic" }]