fix: support sklearn BaseEstimators (#407)

The transformers and threshold model inherit BaseEstimator class.

Fixes:
> mlflow saving of models

---------

Signed-off-by: Kushal Batra <i.kushalbatra@gmail.com>

Author: s0nicboOm

numalogic/registry/mlflow_registry.py

@@ -21,9 +21,9 @@
 from mlflow.exceptions import RestException
 from mlflow.protos.databricks_pb2 import ErrorCode, RESOURCE_DOES_NOT_EXIST
 from mlflow.tracking import MlflowClient
-from sklearn.base import BaseEstimator
 from torch import nn
 
+from numalogic.base import BaseThresholdModel, BaseTransformer
 from numalogic.registry import ArtifactManager, ArtifactData
 from numalogic.registry.artifact import ArtifactCache
 from numalogic.tools.exceptions import ModelVersionError
@@ -104,7 +104,7 @@ def __init__(
     def handler_from_obj(artifact: artifact_t):
         if isinstance(artifact, nn.Module):
             return mlflow.pytorch
-        if isinstance(artifact, BaseEstimator):
+        if isinstance(artifact, (BaseThresholdModel, BaseTransformer)):
             return mlflow.sklearn
         return mlflow.pyfunc
 

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" }]

tests/registry/_mlflow_utils.py

@@ -7,10 +7,11 @@
 from mlflow.entities.model_registry import ModelVersion
 from mlflow.models.model import ModelInfo
 from mlflow.store.entities import PagedList
-from sklearn.ensemble import RandomForestRegressor
 from sklearn.preprocessing import StandardScaler
 from torch import tensor
 
+from numalogic.models.threshold import StdDevThreshold
+
 
 def create_model():
     x = torch.linspace(-math.pi, math.pi, 2000)
@@ -33,8 +34,7 @@ def create_model():
 
 
 def model_sklearn():
-    params = {"n_estimators": 5, "random_state": 42}
-    return RandomForestRegressor(**params)
+    return StdDevThreshold()
 
 
 def mock_log_state_dict(*_, **__):