handlers to single dispatch

Author: isabelizimm

vetiver/__init__.py

@@ -15,5 +15,5 @@
 from .write_docker import *
 from .write_fastapi import * 
 from .handlers._interface import *
-from .handlers.sklearn_vt import *
-from .handlers.pytorch_vt import *
+from .handlers.scikitlearn import *
+from .handlers.pytorch import *

vetiver/handlers/_interface.py

@@ -1,31 +1,100 @@
-from vetiver.handlers import pytorch_vt, sklearn_vt
-import sklearn 
+from typing import Any
+from vetiver.handlers import pytorch, scikitlearn
+from functools import singledispatch
+import sklearn
 
 torch_exists = True
 try:
     import torch
 except ImportError:
     torch_exists = False
 
-def create_translator(model, ptype_data):
+class InvalidModelError(Exception):
+    """
+    Throw an error if `model` is not
+    from scikit-learn or torch
+    """
+
+    def __init__(
+        self,
+        message="The `model` argument must be a scikit-learn or torch model.",
+    ):
+        self.message = message
+        super().__init__(self.message)
+
+CREATE_PTYPE_TPL =  """\
+Failed to create a handler from model of \
+type {_model_type}. If your model is not one of \
+(scikit-learn, torch), you should create and register \
+the handler. Here is a template for such a function: \
+    from pydantic import create_model
+    from vetiver.handlers._interface import create_handler
+    from vetiver.handlers.base import VetiverHandler
+
+    class CustomTemplateHandler(VetiverHandler):
+        def __init__(model, ptype_data):
+            super().__init__(model, ptype_data)
+        
+        def vetiver_create_meta(
+             user: list = None, 
+             version: str = None, 
+             url: str = None, 
+             required_pkgs: list = []):
+        \"""
+        Create metadata for model. This method should include the required 
+        packages necessary to create a prediction.
+        \"""
+        required_pkgs = required_pkgs + ["name_of_modeling_package"]
+        meta = vetiver_meta(user, version, url, required_pkgs)
+        
+        return meta
+
+        def handler_predict(self, input_data, check_ptype):
+            \"""
+            handler_predict should define how to make predictions from your model
+            \"""
+            ...
+
+    @vetiver_create_ptype.register
+    def _(model: {_model_type}, ptype_data):
+        return CustomTemplateHandler(model, ptype_data)
+
+If your datatype is a common type, please consider submitting \
+a pull request.
+"""
+
+@singledispatch
+def create_handler(model, ptype_data):
     """check for model type to handle prediction
 
     Parameters
     ----------
-    model
+    model: object
         Description of parameter `x`.
+    ptype_data : object
+        An object with information (data) whose layout is to be determined.
 
     Returns
     -------
-    pytorch_vt.TorchHandler or sklearn_vt.SKLearnHandler
+    handler
         Handler class for specified model type
+
+    Examples
+    --------
+    >>> import vetiver
+    >>> X, y = vetiver.mock.get_mock_data()
+    >>> model = vetiver.mock.get_mock_model()
+    >>> handler = vetiver.create_handler(model, X)
+    >>> handler.create_description()
+    Scikit-learn <class 'sklearn.dummy.DummyRegressor'> model
     """
-    if torch_exists:
-        if isinstance(model, torch.nn.Module):
-            return pytorch_vt.TorchHandler(model, ptype_data)
+    raise InvalidModelError(message=CREATE_PTYPE_TPL.format(_model_type=type(model)))
 
-    if isinstance(model, sklearn.base.BaseEstimator):
-        return sklearn_vt.SKLearnHandler(model, ptype_data)
+@create_handler.register
+def _(model: sklearn.base.BaseEstimator, ptype_data: Any):
+    return scikitlearn.SKLearnHandler(model, ptype_data)
 
-    else:
-        raise NotImplementedError
+if torch_exists:
+    @create_handler.register
+    def _(model: torch.nn.Module, ptype_data: Any):
+        return pytorch.TorchHandler(model, ptype_data)

vetiver/handlers/base.py

@@ -0,0 +1,79 @@
+from ..ptype import vetiver_create_ptype
+from ..meta import vetiver_meta
+
+class VetiverHandler:
+    """Base handler class for creating VetiverModel of different type.
+
+    Parameters
+    ----------
+    model :
+        a trained model
+    ptype_data:
+        An object with information (data) whose layout is to be determined.
+    """
+
+    def __init__(self, model, ptype_data):
+        self.model = model
+        self.ptype_data = ptype_data
+
+    def create_description(self):
+        """Create description for model"""
+        desc = f"{self.model.__class__} model"
+        return desc
+
+    def vetiver_create_meta(
+        user: list = None,
+        version: str = None,
+        url: str = None,
+        required_pkgs: list = [],
+    ):
+        """Create metadata for sklearn model"""
+        meta = vetiver_meta(user, version, url, required_pkgs)
+
+        return meta
+
+    def ptype(self):
+        """Create data prototype for torch model
+
+        Parameters
+        ----------
+        ptype_data : pd.DataFrame, np.ndarray, or None
+            Training data to create ptype
+
+        Returns
+        -------
+        ptype : pd.DataFrame or None
+            Zero-row DataFrame for storing data types
+        """
+        ptype = vetiver_create_ptype(self.ptype_data)
+        return ptype
+
+    def handler_startup():
+        """Include required packages for prediction
+
+        The `handler_startup` function executes when the API starts. Use this
+        function for tasks like loading packages.
+        """
+        ...
+
+
+    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:
+            Data used to generate prediction
+        check_ptype:
+            If type should be checked against `ptype` or not
+
+        Returns
+        -------
+        prediction
+            Prediction from model
+        """
+        ...

vetiver/handlers/pytorch_vt.py vetiver/handlers/pytorch.pyvetiver/handlers/pytorch_vt.py renamed to vetiver/handlers/pytorch.py

@@ -1,3 +1,4 @@
+from vetiver.handlers.base import VetiverHandler
 from ..meta import vetiver_meta
 from ..ptype import vetiver_create_ptype
 import numpy as np
@@ -9,7 +10,7 @@
     torch_exists = False
 
 
-class TorchHandler:
+class TorchHandler(VetiverHandler):
     """Handler class for creating VetiverModels with torch.
 
     Parameters
@@ -18,8 +19,7 @@ class TorchHandler:
         a trained torch model
     """
     def __init__(self, model, ptype_data):
-        self.model = model
-        self.ptype_data = ptype_data
+        super().__init__(model, ptype_data)
 
     def create_description(self):
         """Create description for torch model

vetiver/handlers/sklearn_vt.py vetiver/handlers/scikitlearn.pyvetiver/handlers/sklearn_vt.py renamed to vetiver/handlers/scikitlearn.py

@@ -1,9 +1,11 @@
+from vetiver.handlers.base import VetiverHandler
 from ..ptype import vetiver_create_ptype
-import pandas as pd
 from ..meta import vetiver_meta
+
+import pandas as pd
 import sklearn
 
-class SKLearnHandler:
+class SKLearnHandler(VetiverHandler):
     """Handler class for creating VetiverModels with sklearn.
 
     Parameters
@@ -13,8 +15,7 @@ class SKLearnHandler:
     """
 
     def __init__(self, model, ptype_data):
-        self.model = model
-        self.ptype_data = ptype_data
+        super().__init__(model, ptype_data)
 
     def create_description(self):
         """Create description for sklearn model

vetiver/tests/test_no_handler.py

@@ -1,11 +1,12 @@
 import pytest 
 from vetiver import VetiverModel, mock
+from vetiver import InvalidModelError
 
 
 def test_not_implemented_error():
     X, y = mock.get_mock_data()
 
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(InvalidModelError):
         VetiverModel(
         model=y,
         ptype_data=X,

vetiver/vetiver_model.py

@@ -1,4 +1,4 @@
-from vetiver.handlers._interface import create_translator
+from vetiver.handlers._interface import create_handler
 
 
 class NoModelAvailableError(Exception):
@@ -44,7 +44,7 @@ def __init__(
         metadata: dict = None,
         **kwargs
     ):
-        translator = create_translator(model, ptype_data)
+        translator = create_handler(model, ptype_data)
 
         self.model = model
         self.ptype = translator.ptype()