use helpers to convert api data to frame

Author: isabelizimm

vetiver/__init__.py

@@ -16,6 +16,7 @@
 from .handlers.torch import TorchHandler  # noqa
 from .handlers.statsmodels import StatsmodelsHandler  # noqa
 from .handlers.xgboost import XGBoostHandler  # noqa
+from .helpers import api_data_to_frame  # noqa
 from .rsconnect import deploy_rsconnect  # noqa
 from .monitor import compute_metrics, pin_metrics, plot_metrics, _rolling_df  # noqa
 from .model_card import model_card  # noqa

vetiver/handlers/base.py

@@ -1,7 +1,5 @@
 from functools import singledispatch
 from contextlib import suppress
-import pandas as pd
-import pydantic
 
 from ..prototype import vetiver_create_prototype
 from ..meta import VetiverMeta
@@ -119,24 +117,6 @@ def handler_startup():
         """
         ...
 
-    def _process_input(self, pred_data) -> pd.DataFrame:
-        """Convert prototype to dataframe data
-
-        Parameters
-        ----------
-        pred_data : pydantic.BaseModel
-            User data from given to API endpoint
-
-        Returns
-        -------
-        pd.DataFrame
-            BaseModel data translated into DataFrame
-        """
-
-        new = _prepare_data(pred_data)
-
-        return new
-
     def handler_predict(self, input_data, check_prototype):
         """Generates method for /predict endpoint in VetiverAPI
 
@@ -169,37 +149,3 @@ def _(model: BaseHandler, prototype_data):
         model.prototype_data = prototype_data
 
     return model
-
-
-@singledispatch
-def _prepare_data(pred_data) -> pd.DataFrame:
-    """Convert prototype to dataframe data
-
-    Parameters
-    ----------
-    pred_data : pydantic.BaseModel
-        User data from given to API endpoint
-
-    Returns
-    -------
-    pd.DataFrame
-        BaseModel data translated into DataFrame
-    """
-
-    raise TypeError("Data should be list, dict, pd.DataFrame")
-
-
-@_prepare_data.register(pydantic.BaseModel)
-@_prepare_data.register(list)
-def _basemodel_list_data(pred_data):
-
-    return pd.DataFrame([dict(s) for s in pred_data])
-
-
-# @_prepare_data.register(dict)
-# def _dict_data(pred_data):
-#     served_data = []
-#     for key, value in pred_data:
-#         served_data.append(value)
-
-#     return served_data

vetiver/handlers/sklearn.py

@@ -2,6 +2,7 @@
 import sklearn
 
 from .base import BaseHandler
+from ..helpers import api_data_to_frame
 
 
 class SKLearnHandler(BaseHandler):
@@ -34,7 +35,7 @@ def handler_predict(self, input_data, check_prototype):
             Prediction from model
         """
         if check_prototype:
-            input_data = self._process_input(input_data)
+            input_data = api_data_to_frame(input_data)
 
         if not check_prototype or isinstance(input_data, pd.DataFrame):
             prediction = self.model.predict(input_data)

vetiver/handlers/statsmodels.py

@@ -1,6 +1,7 @@
 import pandas as pd
 
 from .base import BaseHandler
+from ..helpers import api_data_to_frame
 
 sm_exists = True
 try:
@@ -43,7 +44,8 @@ def handler_predict(self, input_data, check_prototype):
             raise ImportError("Cannot import `statsmodels`")
 
         if check_prototype:
-            input_data = self._process_input(input_data)
+            input_data = api_data_to_frame(input_data)
+
         if isinstance(input_data, (list, pd.DataFrame)):
             prediction = self.model.predict(input_data)
         else:

vetiver/handlers/torch.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from .base import BaseHandler
+from ..helpers import api_data_to_frame
 
 torch_exists = True
 try:
@@ -43,7 +44,7 @@ def handler_predict(self, input_data, check_prototype):
             raise ImportError("Cannot import `torch`.")
 
         if check_prototype:
-            input_data = self._process_input(input_data)
+            input_data = api_data_to_frame(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

@@ -1,6 +1,7 @@
 import pandas as pd
 
 from .base import BaseHandler
+from ..helpers import api_data_to_frame
 
 xgb_exists = True
 try:
@@ -44,7 +45,7 @@ def handler_predict(self, input_data, check_prototype):
             raise ImportError("Cannot import `xgboost`")
 
         if check_prototype:
-            input_data = self._process_input(input_data)
+            input_data = api_data_to_frame(input_data)
 
         if not isinstance(input_data, pd.DataFrame):
             try:

vetiver/helpers.py

@@ -0,0 +1,51 @@
+from functools import singledispatch
+import pandas as pd
+import pydantic
+
+
+@singledispatch
+def api_data_to_frame(pred_data) -> pd.DataFrame:
+    """Convert prototype to dataframe data
+
+    Parameters
+    ----------
+    pred_data : pydantic.BaseModel
+        User data from given to API endpoint
+
+    Returns
+    -------
+    pd.DataFrame
+        BaseModel data translated into DataFrame
+    """
+
+    raise TypeError("Data should be list, pydantic.BaseModel, pd.DataFrame")
+
+
+@api_data_to_frame.register(pydantic.BaseModel)
+@api_data_to_frame.register(list)
+def _(pred_data):
+
+    return pd.DataFrame([dict(s) for s in pred_data])
+
+
+@api_data_to_frame.register(dict)
+def _dict(pred_data):
+    return api_data_to_frame([pred_data])
+
+
+# possible other names
+# prototype_to_frame
+# prototype_to_data
+# prototype_to_dataframe
+# prototype_to_datatype
+# prototype_to_type
+# api_data_to_
+# json_to_
+# server_data_to_
+# transport_data_to_
+# request
+# query
+# transit
+# interchange/exchange
+# transfer
+# through

vetiver/tests/test_add_endpoint.py

@@ -5,6 +5,7 @@
 from fastapi.testclient import TestClient
 
 from vetiver import mock, VetiverModel, VetiverAPI
+from vetiver.helpers import api_data_to_frame
 import vetiver
 
 
@@ -25,12 +26,7 @@ def vetiver_model():
 
 
 def sum_values(x):
-    return pd.DataFrame([dict(s) for s in x]).sum()
-
-
-def sum_dict(x):
-    x = pd.DataFrame(x)
-    return x.sum()
+    return api_data_to_frame(x).sum()
 
 
 @pytest.fixture
@@ -49,7 +45,7 @@ def vetiver_client(vetiver_model):  # With check_prototype=True
 def vetiver_client_check_ptype_false(vetiver_model):  # With check_prototype=False
 
     app = VetiverAPI(vetiver_model, check_prototype=False)
-    app.vetiver_post(sum_dict, "sum")
+    app.vetiver_post(sum_values, "sum")
 
     app.app.root_path = "/sum"
     client = TestClient(app.app)