better testing

Author: isabelizimm

vetiver/__init__.py

@@ -15,5 +15,6 @@
 from .write_docker import *
 from .write_fastapi import * 
 from .handlers._interface import create_handler, InvalidModelError
+from .handlers.base import VetiverHandler
 from .handlers.sklearn import SKLearnHandler
 from .handlers.torch import TorchHandler

vetiver/handlers/_interface.py

@@ -1,6 +1,8 @@
 from vetiver.handlers import torch, sklearn, base
 from functools import singledispatch
 
+from vetiver.ptype import vetiver_create_ptype
+
 class InvalidModelError(Exception):
     """
     Throw an error if `model` is not
@@ -79,10 +81,16 @@ def create_handler(model, ptype_data):
     >>> handler.create_description()
     Scikit-learn <class 'sklearn.dummy.DummyRegressor'> model
     """
+
     raise InvalidModelError(message=CREATE_PTYPE_TPL.format(_model_type=type(model)))
 
 create_handler.register(sklearn.SKLearnHandler.base_class, sklearn.SKLearnHandler)
 
 create_handler.register(torch.TorchHandler.base_class, torch.TorchHandler)
 
-#create_handler.register(base.VetiverHandler, lambda model: model)
+@create_handler.register
+def _(model:base.VetiverHandler, ptype_data):
+    if model.ptype_data is None and ptype_data is not None:
+        model.ptype_data = ptype_data
+    
+    return model

vetiver/handlers/base.py

@@ -1,23 +1,22 @@
 from ..ptype import vetiver_create_ptype
 from ..meta import vetiver_meta
+from abc import ABCMeta
 
-class VetiverHandler:
+class VetiverHandler(metaclass=ABCMeta):
     """Base handler class for creating VetiverModel of different type.
 
     Parameters
     ----------
     model :
         a trained model
-    ptype_data:
+    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 __class__(self):
-        ...
 
     def create_description(self):
         """Create description for model"""

vetiver/tests/test_custom_handler.py

@@ -0,0 +1,60 @@
+import pytest
+import sklearn
+import pydantic
+import pandas as pd
+
+from vetiver import mock, VetiverModel, VetiverHandler
+
+
+class CustomHandler(VetiverHandler):
+    def __init__(self, model, ptype_data):
+        super().__init__(model, ptype_data)
+
+    def handler_predict(self, input_data, check_ptype):
+        if check_ptype == True:
+            if isinstance(input_data, pd.DataFrame):
+                prediction = self.model.predict(input_data)
+            else:
+               prediction = self.model.predict([input_data])
+        else:
+            if not isinstance(input_data, list):
+                input_data = [input_data.split(",")]  # user delimiter ?
+            prediction = self.model.predict(input_data)
+
+        return prediction
+
+
+def test_custom_vetiver_model():
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    custom_handler = CustomHandler(model, X)
+
+    v = VetiverModel(
+        model=custom_handler,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+
+    assert v.description == "A regression model for testing purposes"
+    assert isinstance(v.model, sklearn.dummy.DummyRegressor)
+    assert isinstance(v.ptype.construct(), pydantic.BaseModel)
+
+
+def test_custom_vetiver_model():
+    X, y = mock.get_mock_data()
+    model = mock.get_mock_model().fit(X, y)
+    custom_handler = CustomHandler(model, None)
+
+    v = VetiverModel(
+        model=custom_handler,
+        ptype_data=X,
+        model_name="my_model",
+        versioned=None,
+        description="A regression model for testing purposes",
+    )
+
+    assert v.description == "A regression model for testing purposes"
+    assert isinstance(v.model, sklearn.dummy.DummyRegressor)
+    assert isinstance(v.ptype.construct(), pydantic.BaseModel)

vetiver/vetiver_model.py

@@ -46,7 +46,7 @@ def __init__(
     ):
         translator = create_handler(model, ptype_data)
 
-        self.model = model
+        self.model = translator.model
         self.ptype = translator.ptype()
         self.model_name = model_name
         self.description = description if description else translator.create_description()