adding Torch mixin (#65)

* adding Torch mixin
* typing

---------

Co-authored-by: Aniket Maurya <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

src/litmodels/integrations/mixins.py

@@ -1,11 +1,15 @@
 import pickle
 import tempfile
+import warnings
 from abc import ABC
 from pathlib import Path
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from litmodels import download_model, upload_model
 
+if TYPE_CHECKING:
+    import torch
+
 
 class ModelRegistryMixin(ABC):
     """Mixin for model registry integration."""
@@ -83,3 +87,81 @@ def pull_from_registry(
         if not isinstance(obj, cls):
             raise RuntimeError(f"Unpickled object is not of type {cls.__name__}: {type(obj)}")
         return obj
+
+
+class PyTorchRegistryMixin(ABC):
+    """Mixin for PyTorch model registry integration."""
+
+    def __post_init__(self) -> None:
+        """Post-initialization method to set up the model."""
+        import torch
+
+        # Ensure that the model is in evaluation mode
+        if not isinstance(self, torch.nn.Module):
+            raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(self)}")
+
+    def push_to_registry(
+        self, model_name: Optional[str] = None, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+    ) -> None:
+        """Push the model to the registry.
+
+        Args:
+            model_name: The name of the model. If not use the class name.
+            model_version: The version of the model. If None, the latest version is used.
+            temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+        """
+        import torch
+
+        if not isinstance(self, torch.nn.Module):
+            raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(self)}")
+
+        if model_name is None:
+            model_name = self.__class__.__name__
+        if temp_folder is None:
+            temp_folder = tempfile.gettempdir()
+        torch_path = Path(temp_folder) / f"{model_name}.pth"
+        torch.save(self.state_dict(), torch_path)
+        # todo: dump also object creation arguments so we can dump it and load with model for object instantiation
+        model_registry = f"{model_name}:{model_version}" if model_version else model_name
+        upload_model(name=model_registry, model=torch_path)
+
+    @classmethod
+    def pull_from_registry(
+        cls,
+        model_name: str,
+        model_version: Optional[str] = None,
+        temp_folder: Optional[str] = None,
+        torch_load_kwargs: Optional[dict] = None,
+    ) -> "torch.nn.Module":
+        """Pull the model from the registry.
+
+        Args:
+            model_name: The name of the model.
+            model_version: The version of the model. If None, the latest version is used.
+            temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+            torch_load_kwargs: Additional arguments to pass to `torch.load()`.
+        """
+        import torch
+
+        if temp_folder is None:
+            temp_folder = tempfile.gettempdir()
+        model_registry = f"{model_name}:{model_version}" if model_version else model_name
+        files = download_model(name=model_registry, download_dir=temp_folder)
+        torch_files = [f for f in files if f.endswith(".pth")]
+        if not torch_files:
+            raise RuntimeError(f"No torch file found for model: {model_registry} with {files}")
+        if len(torch_files) > 1:
+            raise RuntimeError(f"Multiple torch files found for model: {model_registry} with {torch_files}")
+        state_dict_path = Path(temp_folder) / torch_files[0]
+        # ignore future warning about changed default
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=FutureWarning)
+            state_dict = torch.load(state_dict_path, **(torch_load_kwargs if torch_load_kwargs else {}))
+
+        # Create a new model instance without calling __init__
+        instance = cls()  # todo: we need to add args used when created dumped model
+        if not isinstance(instance, torch.nn.Module):
+            raise TypeError(f"The model must be a PyTorch `nn.Module` but got: {type(instance)}")
+        # Now load the state dict on the instance
+        instance.load_state_dict(state_dict, strict=True)
+        return instance

tests/integrations/test_mixins.py

@@ -1,6 +1,8 @@
 from unittest import mock
 
-from litmodels.integrations.mixins import PickleRegistryMixin
+import torch
+from litmodels.integrations.mixins import PickleRegistryMixin, PyTorchRegistryMixin
+from torch import nn
 
 
 class DummyModel(PickleRegistryMixin):
@@ -29,3 +31,38 @@ def test_pickle_push_and_pull(mock_download_model, mock_upload_model, tmp_path):
     )
     # Verify that the unpickled instance has the expected value.
     assert loaded_dummy.value == 42
+
+
+class DummyTorchModel(nn.Module, PyTorchRegistryMixin):
+    def __init__(self, input_size=784):
+        super().__init__()
+        self.fc = nn.Linear(input_size, 10)
+
+    def forward(self, x):
+        x = x.view(x.size(0), -1)
+        return self.fc(x)
+
+
+@mock.patch("litmodels.integrations.mixins.upload_model")
+@mock.patch("litmodels.integrations.mixins.download_model")
+def test_pytorch_pull_updated(mock_download_model, mock_upload_model, tmp_path):
+    # Create an instance, push the model and record its forward output.
+    dummy = DummyTorchModel(784)
+    dummy.eval()
+    input_tensor = torch.randn(1, 784)
+    output_before = dummy(input_tensor)
+
+    dummy.push_to_registry(temp_folder=str(tmp_path))
+    expected_path = tmp_path / f"{dummy.__class__.__name__}.pth"
+    mock_upload_model.assert_called_once_with(name="DummyTorchModel", model=expected_path)
+
+    torch.save(dummy.state_dict(), expected_path)
+    # Prepare mocking for pull_from_registry.
+    mock_download_model.return_value = [f"{dummy.__class__.__name__}.pth"]
+    loaded_dummy = DummyTorchModel.pull_from_registry(model_name="DummyTorchModel", temp_folder=str(tmp_path))
+    loaded_dummy.eval()
+    output_after = loaded_dummy(input_tensor)
+
+    assert isinstance(loaded_dummy, DummyTorchModel)
+    # Compare the outputs as a verification.
+    assert torch.allclose(output_before, output_after), "Loaded model output differs from original."