adding load_model for pickled models (#38)

* adding `load_model`
* linter

Author: Borda

README.md

@@ -72,16 +72,13 @@ upload_model(model=model, name=MY_MODEL_NAME)
 ```python
 import os
 import joblib
-from litmodels import download_model
+from litmodels import load_model
 
 # Unique model identifier: <organization>/<teamspace>/<model-name>
 MY_MODEL_NAME = "your_org/your_team/sklearn-svm-model"
 
-# Download the model file from cloud storage
-model_path = download_model(name=MY_MODEL_NAME, download_dir="my_models")
-
-# Load the model for inference using joblib
-model = joblib.load(os.path.join("my_models", model_path[0]))
+# Download and load the model file from cloud storage
+model = load_model(name=MY_MODEL_NAME, download_dir="my_models")
 
 # Example: run inference with the loaded model
 sample_input = [[5.1, 3.5, 1.4, 0.2]]

src/litmodels/__init__.py

@@ -7,6 +7,6 @@
 _PACKAGE_ROOT = os.path.dirname(__file__)
 _PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT)
 
-from litmodels.io import download_model, upload_model
+from litmodels.io import download_model, load_model, upload_model
 
-__all__ = ["download_model", "upload_model"]
+__all__ = ["download_model", "upload_model", "load_model"]

src/litmodels/io/__init__.py

@@ -1,6 +1,6 @@
 """Root package for Input/output."""
 
 from litmodels.io.cloud import download_model_files, upload_model_files
-from litmodels.io.gateway import download_model, upload_model
+from litmodels.io.gateway import download_model, load_model, upload_model
 
-__all__ = ["download_model", "upload_model", "download_model_files", "upload_model_files"]
+__all__ = ["download_model", "upload_model", "download_model_files", "upload_model_files", "load_model"]

src/litmodels/io/gateway.py

@@ -87,3 +87,26 @@ def download_model(
         download_dir=download_dir,
         progress_bar=progress_bar,
     )
+
+
+def load_model(name: str, download_dir: str = ".") -> Any:
+    """Download a model from the model store and load it into memory.
+
+    Args:
+        name: Name of the model to download. Must be in the format 'organization/teamspace/modelname'
+            where entity is either your username or the name of an organization you are part of.
+        download_dir: A path to directory where the model should be downloaded. Defaults
+            to the current working directory.
+
+    Returns:
+        The loaded model.
+    """
+    download_paths = download_model(name=name, download_dir=download_dir)
+    # filter out all Markdown, TXT and RST files
+    download_paths = [p for p in download_paths if Path(p).suffix.lower() not in {".md", ".txt", ".rst"}]
+    if len(download_paths) > 1:
+        raise NotImplementedError("Downloaded model with multiple files is not supported yet.")
+    model_path = Path(os.path.join(download_dir, download_paths[0]))
+    if model_path.suffix.lower() == ".pkl":
+        return joblib.load(model_path)
+    raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.")

tests/test_io_cloud.py

@@ -1,8 +1,9 @@
 import os
 from unittest import mock
 
+import joblib
 import pytest
-from litmodels import download_model, upload_model
+from litmodels import download_model, load_model, upload_model
 from litmodels.io import upload_model_files
 from sklearn import svm
 from torch.nn import Module
@@ -56,3 +57,22 @@ def test_download_model(mock_download_model):
     mock_download_model.assert_called_once_with(
         name="org-name/teamspace/model-name", download_dir="where/to/download", progress_bar=True
     )
+
+
+@mock.patch("litmodels.io.cloud.sdk_download_model")
+def test_load_model(mock_download_model, tmp_path):
+    # create a dummy model file
+    model_file = tmp_path / "dummy_model.pkl"
+    test_data = svm.SVC()
+    joblib.dump(test_data, 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, svm.SVC)