xgboost start

Co-authored-by: Michael Chow <chowmichedu@gmail.com>

Author: isabelizimm

.github/workflows/tests.yml

@@ -32,7 +32,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -e .[dev,torch,statsmodels]
+          python -m pip install -e .[dev,torch,statsmodels,xgboost]
       - name: Run Tests
         run: |
           pytest -m 'not rsc_test' --cov --cov-report xml
@@ -65,8 +65,8 @@ jobs:
         run: |
           pytest vetiver -m 'rsc_test'
 
-  test-no-torch:
-    name: "Test no-torch"
+  test-no-extras:
+    name: "Test no exra ml frameworks"
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2

setup.cfg

@@ -51,3 +51,6 @@ torch =
 
 statsmodels =
     statsmodels
+
+xgboost =
+    xgboost

vetiver/handlers/xgboost.py

@@ -0,0 +1,75 @@
+import pandas as pd
+
+from ..meta import _model_meta
+from .base import BaseHandler
+
+xgb_exists = True
+try:
+    import xgboost
+except ImportError:
+    xgb_exists = False
+
+
+class XGBoostHandler(BaseHandler):
+    """Handler class for creating VetiverModels with statsmodels.
+
+    Parameters
+    ----------
+    model : statsmodels
+        a trained and fit statsmodels model
+    """
+
+    model_class = staticmethod(lambda: xgboost.Booster)
+
+    def __init__(self, model, ptype_data):
+        super().__init__(model, ptype_data)
+
+    def describe(self):
+        """Create description for xgboost model"""
+        desc = f"Statsmodels {self.model.__class__} model."
+        return desc
+
+    def create_meta(
+        user: list = None,
+        version: str = None,
+        url: str = None,
+        required_pkgs: list = [],
+    ):
+        """Create metadata for statsmodel"""
+        required_pkgs = required_pkgs + ["xgboost"]
+        meta = _model_meta(user, version, url, required_pkgs)
+
+        return meta
+
+    def handler_predict(self, input_data, check_ptype):
+        """Generates method for /predict endpoint in VetiverAPI
+
+        The `handler_predict` function executes at each API call. Use this
+        function for calling `predict()` and any other tasks that must be executed
+        at each API call.
+
+        Parameters
+        ----------
+        input_data:
+            Test data
+
+        Returns
+        -------
+        prediction
+            Prediction from model
+        """
+
+        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
+                    )
+
+            prediction = self.model.predict(input_data)
+        else:
+            raise ImportError("Cannot import `xgboost`")
+
+        return prediction

vetiver/tests/test_xgboost.py

@@ -0,0 +1,47 @@
+import pytest
+
+xgb = pytest.importorskip("xgboost", reason="xgboost library not installed")
+
+from vetiver.data import mtcars  # noqa
+from vetiver.handlers.xgboost import XGBoostHandler  # noqa
+
+
+@pytest.fixture
+def fit():
+    # 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"}
+    num_round = 2
+    fit = xgb.train(param, dtrain, num_round)
+
+    return fit
+
+
+@pytest.fixture
+def handler(fit):
+    return XGBoostHandler(fit, None)
+
+
+def test_handler_xgboost_predict_dmatrix(handler):
+    dtest = xgb.DMatrix(mtcars.drop(columns="mpg"))
+    handler.handler_predict(dtest, True)
+
+
+def test_handler_xgboost_predict_df(handler):
+    dtest = mtcars.drop(columns="mpg")
+    handler.handler_predict(dtest, True)
+
+
+@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)
+
+
+@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)