Add Pytorch model abstraction

Author: will-rice

.circleci/config.yml

@@ -19,6 +19,7 @@ jobs:
           pip install --upgrade pip pip-tools
           pip install -r requirements.txt
           pip install tensorflow
+          pip install torch
           python -m spacy download en_core_web_sm
       - save_cache:
           key: build-{{checksum "requirements.txt"}}

spokestack/models/pytorch.py

@@ -0,0 +1,20 @@
+"""Pipeline compatible abstraction for Pytorch jit models."""
+
+import numpy as np
+import torch
+
+
+class PyTorchModel:
+    """Pytorch JIT Model."""
+
+    def __init__(self, model_path: str, device: str = "cpu") -> None:
+        self.model = torch.jit.load(model_path, map_location=device)
+        self.model.eval()
+        self.device = device
+
+    def __call__(self, inputs: np.ndarray) -> np.ndarray:
+        with torch.no_grad():
+            inputs = torch.from_numpy(inputs).to(self.device)
+            out = self.model(inputs)
+
+        return out

tests/models/test_pytorch.py

@@ -0,0 +1,18 @@
+"""
+Tests for Pytorch model base class
+"""
+from unittest import mock
+
+import numpy as np
+
+from spokestack.models.pytorch import PyTorchModel
+
+
+@mock.patch("spokestack.models.pytorch.torch")
+def test_inputs(*args):
+    sample = np.random.rand(1, 128, 80).astype(np.float32)
+    model = PyTorchModel(model_path="torch_model")
+
+    output = model(sample)
+
+    assert output