Merge pull request #50 from has2k1/single-dispatch

Use singledispatch to generate ptypes and remove save_ptype parameter

Author: isabelizimm

examples/coffeeratings.py

@@ -17,7 +17,7 @@
 lr_fit = LinearRegression().fit(X_train, y_train)
 
 # create vetiver model
-v = vetiver.VetiverModel(lr_fit, save_ptype = True, ptype_data=X_train, model_name = "v")
+v = vetiver.VetiverModel(lr_fit, ptype_data=X_train, model_name = "v")
 
 # version model via pin
 from pins import board_folder

vetiver/handlers/_interface.py

@@ -7,7 +7,7 @@
 except ImportError:
     torch_exists = False
 
-def create_translator(model, ptype_data, save_ptype):
+def create_translator(model, ptype_data):
     """check for model type to handle prediction
 
     Parameters
@@ -22,10 +22,10 @@ def create_translator(model, ptype_data, save_ptype):
     """
     if torch_exists:
         if isinstance(model, torch.nn.Module):
-            return pytorch_vt.TorchHandler(model, ptype_data, save_ptype)
+            return pytorch_vt.TorchHandler(model, ptype_data)
 
     if isinstance(model, sklearn.base.BaseEstimator):
-        return sklearn_vt.SKLearnHandler(model, ptype_data, save_ptype)
+        return sklearn_vt.SKLearnHandler(model, ptype_data)
 
     else:
         raise NotImplementedError

vetiver/handlers/pytorch_vt.py

@@ -17,10 +17,9 @@ class TorchHandler:
     model : nn.Module
         a trained torch model
     """
-    def __init__(self, model, ptype_data, save_ptype):
+    def __init__(self, model, ptype_data):
         self.model = model
         self.ptype_data = ptype_data
-        self.save_ptype = save_ptype
 
     def create_description(self):
         """Create description for torch model
@@ -48,14 +47,13 @@ def ptype(self):
         ----------
         ptype_data : pd.DataFrame, np.ndarray, or None
             Training data to create ptype
-        save_ptype : bool
 
         Returns
         -------
         ptype : pd.DataFrame or None
             Zero-row DataFrame for storing data types
         """
-        ptype = vetiver_create_ptype(self.ptype_data, self.save_ptype)
+        ptype = vetiver_create_ptype(self.ptype_data)
 
         return ptype
 
@@ -90,15 +88,10 @@ def handler_predict(self, input_data, check_ptype):
                 prediction = self.model(torch.from_numpy(input_data))
 
             # do not check ptype
-            else:
-                batch = True
-                if not isinstance(input_data, list):
-                    batch = False
-                    input_data = input_data.split(",")  # user delimiter ?
-                input_data = np.array(input_data, dtype=np.array(self.ptype_data).dtype)
-                if not batch:
-                    input_data = input_data.reshape(1, -1)
-                prediction = self.model(torch.from_numpy(input_data))
+            else:    
+                input_data = torch.tensor(input_data)
+                prediction = self.model(input_data)
+
         else:
             raise ImportError("Cannot import `torch`.")
 

vetiver/handlers/sklearn_vt.py

@@ -12,10 +12,9 @@ class SKLearnHandler:
         a trained sklearn model
     """
 
-    def __init__(self, model, ptype_data, save_ptype):
+    def __init__(self, model, ptype_data):
         self.model = model
         self.ptype_data = ptype_data
-        self.save_ptype = save_ptype
 
     def create_description(self):
         """Create description for sklearn model
@@ -42,14 +41,13 @@ def ptype(self):
         ----------
         ptype_data : pd.DataFrame, np.ndarray, or None
             Training data to create ptype
-        save_ptype : bool
 
         Returns
         -------
         ptype : pd.DataFrame or None
             Zero-row DataFrame for storing data types
         """
-        ptype = vetiver_create_ptype(self.ptype_data, self.save_ptype)
+        ptype = vetiver_create_ptype(self.ptype_data)
         return ptype
 
     def handler_startup():

vetiver/pin_read_write.py

@@ -29,7 +29,6 @@ def vetiver_pin_write(board, model: VetiverModel, versioned: bool=True):
         type = "joblib",
         description = model.description,
         metadata = {"required_pkgs": model.metadata.get("required_pkgs"),
-                    "save_ptype": model.save_ptype,
                     "ptype": None if model.ptype == None else model.ptype().json()},
         versioned=versioned
     )
@@ -79,7 +78,6 @@ def vetiver_pin_read(board, name: str, version: str = None) -> VetiverModel:
              url = meta.user.get("url"), # None all the time, besides Connect
              required_pkgs = meta.user.get("required_pkgs")
         ),
-        save_ptype=meta.user.get("save_ptype"),
         ptype_data = json.loads(meta.user.get("ptype")) if meta.user.get("ptype") else None,
         versioned = True
         )

vetiver/ptype.py

@@ -1,7 +1,15 @@
+from functools import singledispatch
+try:
+    from types import NoneType
+except ImportError:
+    # python < 3.10
+    NoneType = type(None)
+
 import pandas as pd
 import numpy as np
 from pydantic import BaseModel, create_model
 
+
 class NoAvailablePTypeError(Exception):
     """
     Throw an error if we cannot create
@@ -18,72 +26,173 @@ def __init__(
 
 class InvalidPTypeError(Exception):
     """
-    Throw an error if `save_ptype` is not
-    True, False, or data.frame
+    Throw an error if ptype cannot be recognised
     """
 
     def __init__(
         self,
-        message="The `ptype_data` argument must be a pandas.DataFrame, a pydantic BaseModel,  np.ndarray, or `save_ptype` must be FALSE.",
+        message="`ptype_data` must be a pd.DataFrame, a pydantic BaseModel or np.ndarray",
     ):
         self.message = message
         super().__init__(self.message)
 
 
-def vetiver_create_ptype(ptype_data, save_ptype: bool):
+CREATE_PTYPE_TPL =  """\
+Failed to create a data prototype (ptype) from data of \
+type {_data_type}. If your datatype is not one of \
+(pd.DataFrame, pydantic.BaseModel, np.ndarry, dict), \
+you should write a function to create the ptype. Here is \
+a template for such a function: \
+
+    from pydantic import create_model
+    from vetiver.ptype import vetiver_create_ptype
+
+    @vetiver_create_ptype.register
+    def _(data: {_data_type}):
+        data_dict = ... # convert data to a dictionary
+        ptype = create_model("ptype", **data_dict)
+        return ptype
+
+If your datatype is a common type, please consider submitting \
+a pull request.
+"""
+
+@singledispatch
+def vetiver_create_ptype(data):
     """Create zero row structure to save data types
+
     Parameters
     ----------
-    ptype_data :
-        Data that represents what
-    save_ptype : bool
-        Whether or not ptype should be created
+    data : object
+        An object with information (data) whose layout is to be determined.
 
     Returns
     -------
-    ptype
+    ptype : pydantic.main.BaseModel
         Data prototype
 
     """
-    ptype = None
+    raise InvalidPTypeError(
+        message=CREATE_PTYPE_TPL.format(_data_type=type(data))
+    )
 
-    if save_ptype == False:
-        pass
-    elif save_ptype == True:
-        try:
-            if isinstance(ptype_data, np.ndarray):
-                ptype = _array_to_ptype(ptype_data[1])
-            elif isinstance(ptype_data, dict):
-                ptype = _dict_to_ptype(ptype_data)
-            elif isinstance(ptype_data.construct(), BaseModel):
-                ptype = ptype_data
-        except AttributeError:  # cannot construct basemodel
-            if isinstance(ptype_data, pd.DataFrame):
-                ptype = _df_to_ptype(ptype_data.iloc[1, :])
-    else:
-        raise InvalidPTypeError
 
+@vetiver_create_ptype.register
+def _(data: pd.DataFrame):
+    """
+    Create ptype for a pandas dataframe
+
+    Parameters
+    ----------
+    data : DataFrame
+        Pandas dataframe
+
+    Examples
+    --------
+    >>> from pydantic import BaseModel
+    >>> df = pd.DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
+    >>> prototype = vetiver_create_ptype(df)
+    >>> issubclass(prototype, BaseModel)
+    True
+    >>> prototype()
+    ptype(x=1, y=4)
+
+    The data prototype created for the dataframe is equivalent to:
+
+    >>> class another_prototype(BaseModel):
+    ...     class Config:
+    ...         title = 'ptype'
+    ...     x: int = 1
+    ...     y: int = 4
+
+    >>> another_prototype()
+    another_prototype(x=1, y=4)
+    >>> another_prototype() == prototype()
+    True
+
+    Changing the title using `class Config` ensures that the
+    also json/schemas match.
+
+    >>> another_prototype.schema() == prototype.schema()
+    True
+    """
+    dict_data = data.iloc[0, :].to_dict()
+    ptype = create_model("ptype", **dict_data)
     return ptype
 
 
-def _df_to_ptype(train_data):
+@vetiver_create_ptype.register
+def _(data: np.ndarray):
+    """
+    Create ptype for a numpy array
 
-    dict_data = train_data.to_dict()
-    ptype = create_model("ptype", **dict_data)
+    Parameters
+    ----------
+    data : ndarray
+        2-Dimensional numpy array
+
+    Examples
+    --------
+    >>> arr = np.array([[1, 4], [2, 5], [3, 6]])
+    >>> prototype = vetiver_create_ptype(arr)
+    >>> prototype()
+    ptype(0=1, 1=4)
+
+    >>> arr2 = np.array([[1, 'a'], [2, 'b'], [3, 'c']], dtype=object)
+    >>> prototype2 = vetiver_create_ptype(arr2)
+    >>> prototype2()
+    ptype(0=1, 1='a')
+    """
+    def _item(value):
+        # pydantic needs python objects. .item() converts a numpy
+        # scalar type to a python equivalent, and if the ndarray
+        # is dtype=object, it may have python objects
+        try:
+            return value.item()
+        except AttributeError:
+            return value
 
+    dict_data = dict(enumerate(data[0], 0))
+    # pydantic requires strings as indicies
+    dict_data = {f"{key}": _item(value) for key, value in dict_data.items()}
+    ptype = create_model("ptype", **dict_data)
     return ptype
 
 
-def _array_to_ptype(train_data):
-    dict_data = dict(enumerate(train_data, 0))
+@vetiver_create_ptype.register
+def _(data: dict):
+    """
+    Create ptype for a dict
 
-    # pydantic requires strings as indicies
-    dict_data = {str(key): value.item() for key, value in dict_data.items()}
-    ptype = create_model("ptype", **dict_data)
+    Parameters
+    ----------
+    data : dict
+        Dictionary
+    """
+    return create_model("ptype", **data)
 
-    return ptype
 
+@vetiver_create_ptype.register
+def _(data: BaseModel):
+    """
+    Create ptype for a pydantic BaseModel object
 
-def _dict_to_ptype(train_data):
+    Parameters
+    ----------
+    data : pydantic.BaseModel
+        Pydantic BaseModel
+    """
+    return data
 
-    return create_model("ptype",**train_data)
+
+@vetiver_create_ptype.register
+def _(data: NoneType):
+    """
+    Create ptype for None
+
+    Parameters
+    ----------
+    data : None
+        None
+    """
+    return None