Merge pull request #24 from isabelizimm/dev-predict-return-df

return dataframes in `predict()`

Author: isabelizimm

vetiver/server.py

@@ -2,6 +2,7 @@
 from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.openapi.utils import get_openapi
+from fastapi import testclient
 
 import uvicorn
 import requests
@@ -169,7 +170,7 @@ def _custom_openapi(self):
         self.app.openapi_schema = openapi_schema
         return self.app.openapi_schema
 
-def predict(endpoint, data: dict, **kw):
+def predict(endpoint, data: Union[dict, pd.DataFrame, pd.Series], **kw):
     """Make a prediction from model endpoint
 
     Parameters
@@ -184,13 +185,35 @@ def predict(endpoint, data: dict, **kw):
     dict
         Key: endpoint_name Value: Output of endpoint_fx, in list format
     """
+    if isinstance(endpoint, testclient.TestClient):
+        requester = endpoint
+        endpoint = "/predict/"
+    else: 
+        requester = requests
+
+    # TO DO: arrow format 
+
     if isinstance(data, pd.DataFrame):
-        data = data.to_json(orient="records")
-        response = requests.post(endpoint, data=data, **kw)
+        data_json = data.to_json(orient="records")
+        response = requester.post(endpoint, data=data_json, **kw)
+    elif isinstance(data, pd.Series):
+        data_dict = data.to_json()
+        response = requester.post(endpoint, data=data_dict, **kw)
+    elif isinstance(data, dict):
+        response = requester.post(endpoint, json=data, **kw)
     else:
-        response = requests.post(endpoint, json=data, **kw)
-
-    return response.json()
+        try:
+            response = requester.post(endpoint, json=data, **kw)
+        except:
+            raise TypeError(f"Predict expects a DataFrame or dict. Given type is {type(data)}")
+
+    response_df = pd.DataFrame.from_dict(response.json())
+    
+    if isinstance(response_df.iloc[0,0], dict):
+        if 'type_error.dict' in response_df.iloc[0,0].values():
+            raise TypeError(f"Predict expects a DataFrame or dict. Given type is {type(data)}")
+
+    return response_df
 
 
 def _prepare_data(pred_data):

vetiver/tests/test_predict.py

@@ -0,0 +1,118 @@
+import pytest
+
+import numpy as np
+import pandas as pd
+from fastapi.testclient import TestClient
+
+from vetiver import mock, VetiverModel, VetiverAPI
+from vetiver.server import predict
+
+
+def test_predict_sklearn_dict_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=True)
+    client = TestClient(app.app)
+    data = {"B": 0, "C": 0, "D": 0}
+    
+    response = predict(endpoint=client, data=data)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+    assert len(response) == 1
+
+
+def test_predict_sklearn_no_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=False)
+    client = TestClient(app.app)
+    
+    response = predict(endpoint=client, data=X)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+    assert len(response) == 100
+
+
+def test_predict_sklearn_df_check_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=True)
+    client = TestClient(app.app)
+    
+    response = predict(endpoint=client, data=X)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+    assert len(response) == 100
+
+
+def test_predict_sklearn_series_check_ptype():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    ser = pd.Series(data=[0,0,0])
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=True)
+    client = TestClient(app.app)
+    
+    response = predict(endpoint=client, data=ser)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.iloc[0,0] == 44.47
+    assert len(response) == 1
+
+
+def test_predict_sklearn_type_error():
+    np.random.seed(500)
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    v = VetiverModel(
+        model=model,
+        save_ptype=True,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+    app = VetiverAPI(v, check_ptype=True)
+    client = TestClient(app.app)
+    data = (0,0)
+    
+    with pytest.raises(TypeError):
+        predict(endpoint=client, data=data)