Support sklearn.tree.DecisionTreeRegressor

2024-06-20 03:29:54 -04:00 · 2024-06-20 03:29:54 -04:00 · 17f6a99fa2
commit 17f6a99fa2
parent a6ed3d96e2
2 changed files with 6 additions and 1 deletions
--- a/mlModelSaver/init.py
+++ b/mlModelSaver/init.py
@ -51,6 +51,10 @@ supportedModels = {
        "supported": True,
        "normalPredictorFunction": "predict_proba"
    },
+    "sklearn.tree.DecisionTreeRegressor": {
+        "supported": True,
+        "normalPredictorFunction": "predict"
+    },
 }

 supportedDataType = {
--- a/pytests/test_mlModelSaver.py
+++ b/pytests/test_mlModelSaver.py
@ -27,7 +27,8 @@ def test_ensureCLassInstance():
        'sm.OLS',
        'sm.Logit',
        'sklearn.neighbors.KNeighborsClassifier',
-        'sklearn.tree.DecisionTreeClassifier'
+        'sklearn.tree.DecisionTreeClassifier',
+        'sklearn.tree.DecisionTreeRegressor',
    ]