batch for torch

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/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,10 +69,11 @@ 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)
@@ -82,6 +83,7 @@ def test_torch_predict_ptype_batch():
     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():
@@ -98,7 +100,7 @@ def test_torch_predict_ptype_error():
 
 
 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)
@@ -107,6 +109,20 @@ def test_torch_predict_no_ptype():
     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"], ["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():