move data handling to handler

Author: isabelizimm

Makefile

@@ -57,6 +57,9 @@ lint:
 test: clean-test
 	pytest -m 'not rsc_test and not docker'
 
+test-pdb: clean-test
+	pytest -m 'not rsc_test and not docker' --pdb
+
 test-rsc: clean-test
 	pytest
 

vetiver/handlers/base.py

@@ -1,6 +1,7 @@
-from vetiver.handlers import base
+from typing import List
 from functools import singledispatch
 from contextlib import suppress
+import pandas as pd
 
 from ..prototype import vetiver_create_prototype
 from ..meta import VetiverMeta
@@ -118,6 +119,30 @@ def handler_startup():
         """
         ...
 
+    def _prepare_data(self, pred_data) -> pd.DataFrame:
+        """Convert prototype to dataframe data
+
+        Parameters
+        ----------
+        prototype_data : pd.DataFrame, np.ndarray, or None
+            Training data to create prototype
+
+        Returns
+        -------
+        prototype : pd.DataFrame or None
+            Zero-row DataFrame for storing data types
+        """        
+        if isinstance(pred_data, List):
+            columns = pred_data[0].dict().keys()
+            data = [line.dict() for line in pred_data]
+            served_data = pd.DataFrame(data, columns=columns)
+        else:
+           served_data = []
+           for key, value in pred_data:
+                served_data.append(value)
+
+        return served_data
+
     def handler_predict(self, input_data, check_prototype):
         """Generates method for /predict endpoint in VetiverAPI
 
@@ -145,7 +170,7 @@ def handler_predict(self, input_data, check_prototype):
 
 
 @create_handler.register
-def _(model: base.BaseHandler, prototype_data):
+def _(model: BaseHandler, prototype_data):
     if model.prototype_data is None and prototype_data is not None:
         model.prototype_data = prototype_data
 

vetiver/handlers/sklearn.py

@@ -33,6 +33,8 @@ def handler_predict(self, input_data, check_prototype):
         prediction
             Prediction from model
         """
+        if check_prototype:
+            input_data = self._prepare_data(input_data)
 
         if not check_prototype or isinstance(input_data, pd.DataFrame):
             prediction = self.model.predict(input_data)

vetiver/handlers/statsmodels.py

@@ -41,7 +41,9 @@ def handler_predict(self, input_data, check_prototype):
         """
         if not sm_exists:
             raise ImportError("Cannot import `statsmodels`")
-
+        
+        if check_prototype:
+            input_data = self._prepare_data(input_data)
         if isinstance(input_data, (list, pd.DataFrame)):
             prediction = self.model.predict(input_data)
         else:

vetiver/handlers/torch.py

@@ -41,7 +41,9 @@ def handler_predict(self, input_data, check_prototype):
         """
         if not torch_exists:
             raise ImportError("Cannot import `torch`.")
+        
         if check_prototype:
+            input_data = self._prepare_data(input_data)
             input_data = np.array(input_data, dtype=np.array(self.prototype_data).dtype)
             prediction = self.model(torch.from_numpy(input_data))
 

vetiver/handlers/xgboost.py

@@ -43,6 +43,9 @@ def handler_predict(self, input_data, check_prototype):
         if not xgb_exists:
             raise ImportError("Cannot import `xgboost`")
 
+        if check_prototype:
+            input_data = self._prepare_data(input_data)
+            
         if not isinstance(input_data, pd.DataFrame):
             try:
                 input_data = pd.DataFrame(input_data)

vetiver/server.py

@@ -171,12 +171,12 @@ async def custom_endpoint(
                 input_data: Union[self.model.prototype, List[self.model.prototype]]
             ):
 
-                if isinstance(input_data, List):
-                    served_data = _batch_data(input_data)
-                else:
-                    served_data = _prepare_data(input_data)
+                # if isinstance(input_data, List):
+                #     served_data = _batch_data(input_data)
+                # else:
+                #     served_data = _prepare_data(input_data)
 
-                new = endpoint_fx(served_data, **kw)
+                new = endpoint_fx(input_data, **kw)
                 return {endpoint_name: new.tolist()}
 
         else:
@@ -287,22 +287,6 @@ def predict(endpoint, data: Union[dict, pd.DataFrame, pd.Series], **kw) -> pd.Da
     return response_df
 
 
-def _prepare_data(pred_data: Dict[str, Any]) -> List[Any]:
-    served_data = []
-    for key, value in pred_data:
-        served_data.append(value)
-    return served_data
-
-
-def _batch_data(pred_data: List[Any]) -> pd.DataFrame:
-    columns = pred_data[0].dict().keys()
-
-    data = [line.dict() for line in pred_data]
-
-    served_data = pd.DataFrame(data, columns=columns)
-    return served_data
-
-
 def vetiver_endpoint(url: str = "http://127.0.0.1:8000/predict") -> str:
     """Wrap url where VetiverModel will be deployed
 

vetiver/tests/test_add_endpoint.py

@@ -25,7 +25,7 @@ def vetiver_model():
 
 
 def sum_values(x):
-    return x.sum()
+    return pd.DataFrame([dict(s) for s in x]).sum()
 
 
 def sum_dict(x):