Merge pull request #325 from NeroHin/master

alejandroschuler · web-flow · commit acf6b270d19d · 2024-10-28T17:14:07.000-07:00
feat: Add estimator type with reg and clf for NGBClassifier and NGBRegressor
diff --git a/ngboost/api.py b/ngboost/api.py
@@ -97,6 +97,8 @@ def __init__(
             early_stopping_rounds,
         )
 
+        self._estimator_type = "regressor"
+
     def __getstate__(self):
         state = super().__getstate__()
         # Remove the unpicklable entries.
@@ -172,6 +174,7 @@ def __init__(
             tol,
             random_state,
         )
+        self._estimator_type = "classifier"
 
     def predict_proba(self, X, max_iter=None):
         """

Original file line number	Diff line number	Diff line change
`@@ -97,6 +97,8 @@ def __init__(`
`97`	`97`	`early_stopping_rounds,`
`98`	`98`	`)`
`99`	`99`
	`100`	`+ self._estimator_type = "regressor"`
	`101`	`+`
`100`	`102`	`def __getstate__(self):`
`101`	`103`	`state = super().__getstate__()`
`102`	`104`	`# Remove the unpicklable entries.`
`@@ -172,6 +174,7 @@ def __init__(`
`172`	`174`	`tol,`
`173`	`175`	`random_state,`
`174`	`176`	`)`
	`177`	`+ self._estimator_type = "classifier"`
`175`	`178`
`176`	`179`	`def predict_proba(self, X, max_iter=None):`
`177`	`180`	`"""`