download_model

Author: Borda

examples/demo-upload-download.py

@@ -16,7 +16,7 @@
     )
 
     # Download the model checkpoint
-    model_path = litmodels.download_model_files("jirka/kaggle/boring-model", download_dir="./my-models")
+    model_path = litmodels.download_model("jirka/kaggle/boring-model", download_dir="./my-models")
     print(f"Model downloaded to {model_path}")
 
     # Load the model checkpoint

examples/train-resume.py

@@ -1,14 +1,14 @@
 import torch.utils.data as data
 import torchvision as tv
 from lightning import Trainer
-from litmodels import download_model_files
+from litmodels import download_model
 from sample_model import LitAutoEncoder
 
 if __name__ == "__main__":
     dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
     train, val = data.random_split(dataset, [55000, 5000])
 
-    model_path = download_model_files(name="jirka/kaggle/lit-auto-encoder-simple", download_dir="my_models")
+    model_path = download_model(name="jirka/kaggle/lit-auto-encoder-simple", download_dir="my_models")
     print(f"model: {model_path}")
     # autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint_path=model_path)
 

src/litmodels/__init__.py

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

src/litmodels/cloud_io.py

@@ -132,7 +132,7 @@ def upload_model_files(
     )
 
 
-def download_model_files(
+def download_model(
     name: str,
     download_dir: str = ".",
     progress_bar: bool = True,

tests/test_cloud_io.py

@@ -2,7 +2,7 @@
 from unittest import mock
 
 import pytest
-from litmodels.cloud_io import download_model_files, upload_model, upload_model_files
+from litmodels.cloud_io import download_model, upload_model, upload_model_files
 from torch.nn import Module
 
 
@@ -11,7 +11,7 @@ def test_wrong_model_name(name):
     with pytest.raises(ValueError, match=r".*organization/teamspace/model.*"):
         upload_model_files(path="path/to/checkpoint", name=name)
     with pytest.raises(ValueError, match=r".*organization/teamspace/model.*"):
-        download_model_files(name=name)
+        download_model(name=name)
 
 
 @pytest.mark.parametrize(
@@ -48,7 +48,7 @@ def test_download_model(mocker):
     ts_mock = mock.MagicMock()
     mocker.patch("litmodels.cloud_io._get_teamspace", return_value=ts_mock)
     # The lit-logger function is just a wrapper around the SDK function
-    download_model_files(
+    download_model(
         name="org-name/teamspace/model-name",
         download_dir="where/to/download",
     )