initial statsmodels support

Author: isabelizimm

vetiver/__init__.py

@@ -14,6 +14,7 @@
 from .handlers.base import BaseHandler, create_handler, InvalidModelError  # noqa
 from .handlers.sklearn import SKLearnHandler  # noqa
 from .handlers.torch import TorchHandler  # noqa
+from .handlers.statsmodels import StatsmodelsHandler  # noqa
 from .rsconnect import deploy_rsconnect  # noqa
 from .monitor import compute_metrics, pin_metrics, plot_metrics, _rolling_df  # noqa
 

vetiver/handlers/statsmodels.py

@@ -0,0 +1,77 @@
+import pandas as pd
+
+from ..meta import _model_meta
+from .base import BaseHandler
+
+sm_exists = True
+try:
+    import statsmodels
+except ImportError:
+    sm_exists = False
+
+
+class StatsmodelsHandler(BaseHandler):
+    """Handler class for creating VetiverModels with sklearn.
+
+    Parameters
+    ----------
+    model : statsmodels
+        a trained sklearn model
+    """
+
+    model_class = staticmethod(
+        lambda: statsmodels.regression.linear_model.RegressionResultsWrapper
+    )
+
+    def __init__(self, model, ptype_data):
+        super().__init__(model, ptype_data)
+
+    def describe(self):
+        """Create description for sklearn model"""
+        desc = f"Statsmodels {self.model.__class__} model."
+        return desc
+
+    def construct_meta(
+        user: list = None,
+        version: str = None,
+        url: str = None,
+        required_pkgs: list = [],
+    ):
+        """Create metadata for sklearn model"""
+        required_pkgs = required_pkgs + ["statsmodels"]
+        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 check_ptype:
+            if isinstance(input_data, pd.DataFrame):
+                prediction = self.model.predict(input_data)
+            else:
+                prediction = self.model.predict([input_data])
+
+        # do not check ptype
+        else:
+            if not isinstance(input_data, list):
+                input_data = [input_data.split(",")]  # user delimiter ?
+
+            prediction = self.model.predict(input_data)
+
+        return prediction

vetiver/tests/test_statsmodels.py

@@ -0,0 +1,128 @@
+# import pytest
+
+# sm = pytest.importorskip("statsmodels", reason="statsmodels library not installed")
+
+# import numpy as np  # noqa
+# from fastapi.testclient import TestClient  # noqa
+
+# from vetiver.vetiver_model import VetiverModel  # noqa
+# from vetiver import VetiverAPI  # noqa
+
+
+# def _build_sm():
+
+#     input_size = 1
+#     output_size = 1
+
+#     x_train = np.array(
+#         [
+#             [3.3],
+#             [4.4],
+#             [5.5],
+#             [6.71],
+#             [6.93],
+#             [4.168],
+#             [9.779],
+#             [6.182],
+#             [7.59],
+#             [2.167],
+#             [7.042],
+#             [10.791],
+#             [5.313],
+#             [7.997],
+#             [3.1],
+#         ],
+#         dtype=np.float32,
+#     )
+
+#     torch_model = sm.nn.Linear(input_size, output_size)
+#     return x_train, torch_model
+
+
+# def test_vetiver_build():
+
+#     x_train, torch_model = _build_sm()
+
+#     vt2 = VetiverModel(
+#         model=torch_model,
+#         ptype_data=x_train,
+#         model_name="torch",
+#         versioned=None,
+#         description=None,
+#         metadata=None,
+#     )
+
+#     assert vt2.model == torch_model
+
+
+# def test_sm_predict_ptype():
+#     torch.manual_seed(3)
+#     x_train, torch_model = _build_sm()
+#     v = VetiverModel(torch_model, model_name="torch", ptype_data=x_train)
+#     v_api = VetiverAPI(v)
+
+#     client = TestClient(v_api.app)
+#     data = {"0": 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_sm_predict_ptype_batch():
+
+#     x_train, torch_model = _build_sm()
+#     v = VetiverModel(torch_model, model_name="torch", 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_sm_predict_ptype_error():
+
+#     x_train, torch_model = _build_sm()
+#     v = VetiverModel(torch_model, model_name="torch", ptype_data=x_train)
+#     v_api = VetiverAPI(v)
+
+#     client = TestClient(v_api.app)
+#     data = {"0": "bad"}
+#     response = client.post("/predict", json=data)
+
+#     assert response.status_code == 422, response.text  # value is not a valid float
+
+
+# def test_sm_predict_no_ptype_batch():
+
+#     x_train, torch_model = _build_sm()
+#     v = VetiverModel(torch_model, model_name="torch")
+#     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_sm_predict_no_ptype():
+
+#     x_train, torch_model = _build_sm()
+#     v = VetiverModel(torch_model, model_name="torch")
+#     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_pin_sm():