handle dict, df inputs

Author: isabelizimm

vetiver/handlers/xgboost.py

@@ -61,12 +61,12 @@ def handler_predict(self, input_data, check_ptype):
 
         if xgb_exists:
             if not isinstance(input_data, xgboost.DMatrix):
-                if isinstance(input_data, pd.DataFrame):
-                    input_data = xgboost.DMatrix(input_data)
-                else:
-                    input_data = xgboost.DMatrix(
-                        input_data, label=self.model.feature_names
-                    )
+                if not isinstance(input_data, pd.DataFrame):
+                    try:
+                        input_data = pd.DataFrame(input_data)
+                    except ValueError:
+                        raise (f"Expected a dict or DataFrame, got {type(input_data)}")
+                input_data = xgboost.DMatrix(input_data)
 
             prediction = self.model.predict(input_data)
         else:

vetiver/server.py

@@ -227,6 +227,7 @@ def predict(endpoint, data: Union[dict, pd.DataFrame, pd.Series], **kw):
     elif isinstance(data, dict):
         response = requester.post(endpoint, json=data, **kw)
     else:
+        # TODO: Check in on JSON serialization of DMatrix for XGBoost
         try:
             response = requester.post(endpoint, json=data, **kw)
         except TypeError:

vetiver/tests/test_xgboost.py

@@ -4,44 +4,69 @@
 
 from vetiver.data import mtcars  # noqa
 from vetiver.handlers.xgboost import XGBoostHandler  # noqa
+import numpy as np  # noqa
+from fastapi.testclient import TestClient  # noqa
+
+import vetiver  # noqa
 
 
 @pytest.fixture
-def fit():
+def build_xgb():
     # read in data
     dtrain = xgb.DMatrix(mtcars.drop(columns="mpg"), label=mtcars["mpg"])
     # specify parameters via map
-    param = {"max_depth": 2, "eta": 1, "objective": "reg:squarederror"}
+    param = {
+        "max_depth": 2,
+        "eta": 1,
+        "objective": "reg:squarederror",
+        "random_state": 0,
+    }
     num_round = 2
     fit = xgb.train(param, dtrain, num_round)
 
-    return fit
+    return vetiver.VetiverModel(fit, "xgb", mtcars.drop(columns="mpg"))
 
 
-@pytest.fixture
-def handler(fit):
-    return XGBoostHandler(fit, None)
+def test_vetiver_build(build_xgb):
+    api = vetiver.VetiverAPI(build_xgb)
+    client = TestClient(api.app)
+    data = mtcars.head(1).drop(columns="mpg")
+
+    response = vetiver.predict(endpoint=client, data=data)
+
+    assert response.iloc[0, 0] == 21.064373016357422
+    assert len(response) == 1
+
+
+def test_batch(build_xgb):
+    api = vetiver.VetiverAPI(build_xgb)
+    client = TestClient(api.app)
+    data = mtcars.head(3).drop(columns="mpg")
+
+    response = vetiver.predict(endpoint=client, data=data)
 
+    assert response.iloc[0, 0] == 21.064373016357422
+    assert len(response) == 3
 
-def test_handler_xgboost_predict_dmatrix(handler):
-    dtest = xgb.DMatrix(mtcars.drop(columns="mpg"))
-    handler.handler_predict(dtest, True)
 
+def test_no_ptype(build_xgb):
+    api = vetiver.VetiverAPI(build_xgb, check_ptype=False)
+    client = TestClient(api.app)
+    data = mtcars.head(1).drop(columns="mpg")
 
-def test_handler_xgboost_predict_df(handler):
-    dtest = mtcars.drop(columns="mpg")
-    handler.handler_predict(dtest, True)
+    response = vetiver.predict(endpoint=client, data=data)
 
+    assert response.iloc[0, 0] == 21.064373016357422
+    assert len(response) == 1
 
-@pytest.mark.xfail
-def test_handler_xgboost_predict_str(handler):
-    # TODO: prediction from a string")
-    dtest = mtcars.drop(columns="mpg")
-    handler.handler_predict(dtest, False)
 
+def test_serialize(build_xgb):
+    import pins
 
-@pytest.mark.xfail
-def test_handler_xgboost_predict_list(handler):
-    # TODO: prediction from a serialized JSON list
-    row = mtcars.drop(columns="mpg").iloc[0, :].tolist()
-    handler.handler_predict([row], False)
+    board = pins.board_temp(allow_pickle_read=True)
+    vetiver.vetiver_pin_write(board=board, model=build_xgb)
+    assert isinstance(
+        board.pin_read("xgb"),
+        xgb.Booster,
+    )
+    board.pin_delete("xgb")