Merge pull request #13 from isabelizimm/dev-ptype

support batch processing for `sklearn` and `pytorch`

Author: isabelizimm

vetiver/handlers/pytorch_vt.py

@@ -72,10 +72,16 @@ def handler_predict(self, input_data, check_ptype):
         if check_ptype == True:
             input_data = np.array(input_data, dtype=np.array(self.ptype_data).dtype)
             prediction = self.model(torch.from_numpy(input_data))
+        
+        # do not check ptype
         else:
-            input_data = input_data.split(",")  # user delimiter ?
+            batch = True
+            if not isinstance(input_data, list):
+                batch = False
+                input_data = input_data.split(",")  # user delimiter ?
             input_data = np.array(input_data, dtype=np.array(self.ptype_data).dtype)
-            reshape_data = input_data.reshape(1, -1)
-            prediction = self.model(torch.from_numpy(reshape_data))
+            if not batch:
+                input_data = input_data.reshape(1, -1)
+            prediction = self.model(torch.from_numpy(input_data))
 
         return prediction

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

@@ -1,14 +1,14 @@
+from cgitb import reset
 import pytest
 
 from vetiver.vetiver_model import VetiverModel
 from vetiver import VetiverAPI
 from fastapi.testclient import TestClient
 
+import torch
 import torch.nn as nn
 import numpy as np
 
-np.random.seed(500)
-
 
 def _build_torch_v():
 
@@ -59,7 +59,7 @@ def test_vetiver_build():
 
 
 def test_torch_predict_ptype():
-
+    torch.manual_seed(3)
     x_train, torch_model = _build_torch_v()
     v = VetiverModel(torch_model, save_ptype=True, ptype_data=x_train)
     v_api = VetiverAPI(v)
@@ -69,6 +69,22 @@ def test_torch_predict_ptype():
     response = client.post("/predict/", json=data)
 
     assert response.status_code == 200, response.text
+    assert response.json() == {"prediction":[-4.060722351074219]}, response.text
+
+
+def test_torch_predict_ptype_batch():
+    torch.manual_seed(3)
+    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
+    assert response.json() == {"prediction":[[-4.060722351074219],[-4.060722351074219]]}, response.text
+
 
 def test_torch_predict_ptype_error():
 
@@ -80,19 +96,33 @@ 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():
+    torch.manual_seed(3)
+    x_train, torch_model = _build_torch_v()
+    v = VetiverModel(torch_model, save_ptype=False, ptype_data=x_train)
+    v_api = VetiverAPI(v, check_ptype=False)
+
+    client = TestClient(v_api.app)
+    data = "3.3"
+    response = client.post("/predict/", json=data)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"prediction":[[-4.060722351074219]]}, response.text
+
 
+def test_torch_predict_no_ptype_batch():
+    torch.manual_seed(3)
     x_train, torch_model = _build_torch_v()
     v = VetiverModel(torch_model, save_ptype=False, ptype_data=x_train)
     v_api = VetiverAPI(v, check_ptype=False)
 
     client = TestClient(v_api.app)
-    data = '3.3'
+    data = [["3.3"], ["3.3"]]
     response = client.post("/predict/", json=data)
     assert response.status_code == 200, response.text
+    assert response.json() == {"prediction":[[-4.060722351074219],[-4.060722351074219]]}, response.text
 
 
 def test_torch_predict_no_ptype_error():
@@ -102,6 +132,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)