updates from review

Author: isabelizimm

vetiver/__init__.py

@@ -14,6 +14,6 @@
 from .meta import *
 from .write_docker import *
 from .write_fastapi import * 
-from .handlers._interface import *
+from .handlers._interface import create_handler
 from .handlers.sklearn import SKLearnHandler
 from .handlers.torch import TorchHandler

vetiver/handlers/_interface.py

@@ -1,5 +1,4 @@
-from typing import Any
-from vetiver.handlers import torch, sklearn
+from vetiver.handlers import torch, sklearn, base
 from functools import singledispatch
 
 class InvalidModelError(Exception):
@@ -20,7 +19,6 @@ def __init__(
 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
 
@@ -86,3 +84,5 @@ def create_handler(model, ptype_data):
 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)

vetiver/handlers/base.py

@@ -16,6 +16,9 @@ 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"""
         desc = f"{self.model.__class__} model"

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

@@ -13,7 +13,8 @@ class SKLearnHandler(VetiverHandler):
     model : sklearn.base.BaseEstimator
         a trained sklearn model
     """
-
+    base_class = sklearn.base.BaseEstimator
+    
     def __init__(self, model, ptype_data):
         super().__init__(model, ptype_data)
 

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

@@ -18,6 +18,7 @@ class TorchHandler(VetiverHandler):
     model : nn.Module
         a trained torch model
     """
+    base_class = torch.nn.Module
     def __init__(self, model, ptype_data):
         super().__init__(model, ptype_data)