load

Author: Borda

src/litmodels/io/gateway.py

@@ -112,4 +112,6 @@ def load_model(name: str, download_dir: str = ".") -> Any:
     model_path = Path(os.path.join(download_dir, download_paths[0]))
     if model_path.suffix.lower() == ".pkl":
         return joblib.load(model_path)
+    if model_path.suffix.lower() == ".pt":
+        return torch.jit.load(model_path)
     raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.")

tests/test_io_cloud.py

@@ -3,6 +3,7 @@
 
 import joblib
 import pytest
+import torch
 import torch.jit as torch_jit
 from litmodels import download_model, load_model, upload_model
 from litmodels.io import upload_model_files
@@ -70,7 +71,7 @@ def test_download_model(mock_download_model):
 
 
 @mock.patch("litmodels.io.cloud.sdk_download_model")
-def test_load_model(mock_download_model, tmp_path):
+def test_load_model_pickle(mock_download_model, tmp_path):
     # create a dummy model file
     model_file = tmp_path / "dummy_model.pkl"
     test_data = svm.SVC()
@@ -86,3 +87,22 @@ def test_load_model(mock_download_model, tmp_path):
         name="org-name/teamspace/model-name", download_dir=str(tmp_path), progress_bar=True
     )
     assert isinstance(model, svm.SVC)
+
+
+@mock.patch("litmodels.io.cloud.sdk_download_model")
+def test_load_model_torch_jit(mock_download_model, tmp_path):
+    # create a dummy model file
+    model_file = tmp_path / "dummy_model.pt"
+    test_data = torch_jit.script(Module())
+    test_data.save(model_file)
+    mock_download_model.return_value = [str(model_file.name)]
+
+    # The lit-logger function is just a wrapper around the SDK function
+    model = load_model(
+        name="org-name/teamspace/model-name",
+        download_dir=str(tmp_path),
+    )
+    mock_download_model.assert_called_once_with(
+        name="org-name/teamspace/model-name", download_dir=str(tmp_path), progress_bar=True
+    )
+    assert isinstance(model, torch.jit.ScriptModule)