batch processing for sklearn

Author: isabelizimm

vetiver/handlers/sklearn_vt.py

@@ -1,5 +1,6 @@
 from ..ptype import _vetiver_create_ptype
-import sklearn
+
+import pandas as pd
 import numpy as np
 
 class SKLearnHandler:
@@ -70,11 +71,16 @@ def handler_predict(self, input_data, check_ptype):
             Prediction from model
         """
         if check_ptype == True:
-            prediction = self.model.predict([input_data])
+            if isinstance(input_data, pd.DataFrame):
+                prediction = self.model.predict(input_data)
+            else:
+               prediction = self.model.predict([input_data])
+
+        # do not check ptype
         else:
-            input_data = input_data.split(",")  # user delimiter ?
-            input_data = np.asarray(input_data)
-            reshape_data = input_data.reshape(1, -1)
-            prediction = self.model.predict(reshape_data)
+            if not isinstance(input_data, list):
+                input_data = [input_data.split(",")]  # user delimiter ?
+
+            prediction = self.model.predict(input_data)
 
         return prediction

vetiver/server.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse
 import uvicorn
-from typing import Callable, Optional
+from typing import Callable, Optional, Union, List
 import requests
 import pandas as pd
 
@@ -81,9 +81,12 @@ async def rapidoc():
         if self.check_ptype == True:
 
             @app.post("/predict/")
-            async def prediction(input_data: self.model.ptype):
-
-                served_data = _prepare_data(input_data)
+            async def prediction(input_data: Union[self.model.ptype, List[self.model.ptype]]):
+                
+                if isinstance(input_data, List):
+                    served_data = _batch_data(input_data)
+                else:
+                    served_data = _prepare_data(input_data)
 
                 y = self.model.handler_predict(served_data, check_ptype=self.check_ptype)
 
@@ -165,6 +168,16 @@ def _prepare_data(pred_data):
         served_data.append(value)
     return served_data
 
+def _batch_data(pred_data):
+    columns = pred_data[0].dict().keys()
+
+    data = [line.dict() for line in pred_data]
+    print(data)
+
+    served_data = pd.DataFrame(data, columns=columns)
+    return served_data
+
+
 
 def vetiver_endpoint(url="http://127.0.0.1:8000/predict"):
     """Wrap url

vetiver/tests/test_pytorch.py

@@ -70,6 +70,20 @@ def test_torch_predict_ptype():
 
     assert response.status_code == 200, response.text
 
+
+def test_torch_predict_ptype_batch():
+
+    x_train, torch_model = _build_torch_v()
+    v = VetiverModel(torch_model, save_ptype=True, ptype_data=x_train)
+    v_api = VetiverAPI(v)
+
+    client = TestClient(v_api.app)
+    data = [{"0": 3.3}, {"0": 3.3}]
+    response = client.post("/predict/", json=data)
+
+    assert response.status_code == 200, response.text
+
+
 def test_torch_predict_ptype_error():
 
     x_train, torch_model = _build_torch_v()
@@ -80,7 +94,7 @@ def test_torch_predict_ptype_error():
     data = {"0": "bad"}
     response = client.post("/predict/", json=data)
 
-    assert response.status_code == 422, response.text # value is not a valid float
+    assert response.status_code == 422, response.text  # value is not a valid float
 
 
 def test_torch_predict_no_ptype():
@@ -90,7 +104,7 @@ def test_torch_predict_no_ptype():
     v_api = VetiverAPI(v, check_ptype=False)
 
     client = TestClient(v_api.app)
-    data = '3.3'
+    data = "3.3"
     response = client.post("/predict/", json=data)
     assert response.status_code == 200, response.text
 
@@ -102,6 +116,6 @@ def test_torch_predict_no_ptype_error():
     v_api = VetiverAPI(v, check_ptype=False)
 
     client = TestClient(v_api.app)
-    data = 'bad'
+    data = "bad"
     with pytest.raises(ValueError):
         client.post("/predict/", json=data)

vetiver/tests/test_sklearn.py

@@ -31,6 +31,14 @@ def test_predict_endpoint_ptype():
     assert response.status_code == 200, response.text
     assert response.json() == {"prediction": [44.47]}, response.json()
 
+def test_predict_endpoint_ptype_batch():
+    np.random.seed(500)
+    client = TestClient(_start_application(save_ptype=True).app)
+    data = [{"B": 0, "C": 0, "D": 0},{"B": 0, "C": 0, "D": 0}]
+    response = client.post("/predict/", json=data)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"prediction": [44.47, 44.47]}, response.json()
+
 
 def test_predict_endpoint_ptype_error():
     np.random.seed(500)
@@ -48,10 +56,17 @@ def test_predict_endpoint_no_ptype():
     assert response.status_code == 200, response.text
     assert response.json() == {"prediction": [44.47]}, response.json()
 
+def test_predict_endpoint_no_ptype_batch():
+    np.random.seed(500)
+    client = TestClient(_start_application(save_ptype=False).app)
+    data = [['0,0,0'],['0,0,0']]
+    response = client.post("/predict/", json=data)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"prediction": [44.47, 44.47]}, response.json()
 
 def test_predict_endpoint_no_ptype_error():
     np.random.seed(500)
     client = TestClient(_start_application(save_ptype=False).app)
-    data = ['hell0',9,32.0]
-    with pytest.raises(AttributeError):
+    data = {'hell0',9,32.0}
+    with pytest.raises(TypeError):
         client.post("/predict/", json=data)