generalize model calls

isabelizimm · isabelizimm · commit 0d089ab4fb71 · 2025-03-20T17:38:31.000-04:00
diff --git a/vetiver/handlers/sklearn.py b/vetiver/handlers/sklearn.py
@@ -46,7 +46,4 @@ def handler_predict(self, input_data, check_prototype: bool, **kw):
             else input_data
         )
 
-        if prediction_type in ["predict_proba", "predict_log_proba"]:
-            return getattr(self.model, prediction_type)(input_data).tolist()
-
-        return self.model.predict(input_data).to_list()
+        return getattr(self.model, prediction_type)(input_data).tolist()

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,4 @@ def handler_predict(self, input_data, check_prototype: bool, **kw):`
`46`	`46`	`else input_data`
`47`	`47`	`)`
`48`	`48`
`49`		`- if prediction_type in ["predict_proba", "predict_log_proba"]:`
`50`		`- return getattr(self.model, prediction_type)(input_data).tolist()`
`51`		`-`
`52`		`- return self.model.predict(input_data).to_list()`
	`49`	`+ return getattr(self.model, prediction_type)(input_data).tolist()`