Merge pull request #29 from isabelizimm/dev-optional-deps

make torch optional dependency

Author: isabelizimm

.github/workflows/tests.yml

@@ -23,7 +23,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -e .[dev]
+          python -m pip install -e .[dev,torch]
       - name: Run Tests
         run: |
           pytest --cov --cov-report xml .

setup.cfg

@@ -31,7 +31,6 @@ install_requires =
     uvicorn
     scikit-learn
     nest-asyncio
-    torch
     requests
     pins @ git+https://github.com/machow/pins-python@f9391b5937b3a13b2a2c841855ca94ab281b2cef
 
@@ -43,3 +42,6 @@ dev =
     sphinx-autodoc-typehints
     sphinx-book-theme
     myst-parser
+
+torch = 
+    torch

vetiver/handlers/_interface.py

@@ -1,6 +1,11 @@
 from vetiver.handlers import pytorch_vt, sklearn_vt
-from torch import nn
-import sklearn
+import sklearn 
+
+torch_exists = True
+try:
+    import torch
+except ImportError:
+    torch_exists = False
 
 def create_translator(model, ptype_data, save_ptype):
     """check for model type to handle prediction
@@ -15,11 +20,11 @@ def create_translator(model, ptype_data, save_ptype):
     pytorch_vt.TorchHandler or sklearn_vt.SKLearnHandler
         Handler class for specified model type
     """
+    if torch_exists:
+        if isinstance(model, torch.nn.Module):
+            return pytorch_vt.TorchHandler(model, ptype_data, save_ptype)
 
-    if isinstance(model, nn.Module):
-        return pytorch_vt.TorchHandler(model, ptype_data, save_ptype)
-
-    elif isinstance(model, sklearn.base.BaseEstimator):
+    if isinstance(model, sklearn.base.BaseEstimator):
         return sklearn_vt.SKLearnHandler(model, ptype_data, save_ptype)
 
     else:

vetiver/handlers/pytorch_vt.py

@@ -2,6 +2,13 @@
 from ..ptype import _vetiver_create_ptype
 import numpy as np
 
+torch_exists = True
+try:
+    import torch
+except ImportError:
+    torch_exists = False
+
+
 class TorchHandler:
     """Handler class for creating VetiverModels with torch.
 
@@ -77,21 +84,22 @@ def handler_predict(self, input_data, check_ptype):
         prediction
             Prediction from model
         """
-        import torch
-        
-        if check_ptype == True:
-            input_data = np.array(input_data, dtype=np.array(self.ptype_data).dtype)
-            prediction = self.model(torch.from_numpy(input_data))
-        
-        # do not check ptype
+        if torch_exists:
+            if check_ptype == True:
+                input_data = np.array(input_data, dtype=np.array(self.ptype_data).dtype)
+                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:
-            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))
+            raise ImportError("Cannot import `torch`.")
 
         return prediction