files

Author: Borda

src/litmodels/io/__init__.py

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

src/litmodels/io/cloud.py

@@ -50,7 +50,7 @@ def _get_teamspace(name: str, organization: str) -> Teamspace:
     return Teamspace(**teamspaces[requested_teamspace])
 
 
-def upload_model_file(
+def upload_model_files(
     name: str,
     path: str,
     progress_bar: bool = True,
@@ -77,7 +77,7 @@ def upload_model_file(
     )
 
 
-def download_model_file(
+def download_model_files(
     name: str,
     download_dir: str = ".",
     progress_bar: bool = True,

src/litmodels/io/gateway.py

@@ -6,7 +6,7 @@
 from lightning_sdk.api.teamspace_api import UploadedModelInfo
 from lightning_utilities import module_available
 
-from litmodels.io.cloud import download_model_file, upload_model_file
+from litmodels.io.cloud import download_model_files, upload_model_files
 
 if module_available("torch"):
     import torch
@@ -49,7 +49,7 @@ def upload_model(
         path = str(model)
     else:
         raise ValueError(f"Unsupported model type {type(model)}")
-    return upload_model_file(
+    return upload_model_files(
         path=path,
         name=name,
         progress_bar=progress_bar,
@@ -74,7 +74,7 @@ def download_model(
     Returns:
         The absolute path to the downloaded model file or folder.
     """
-    return download_model_file(
+    return download_model_files(
         name=name,
         download_dir=download_dir,
         progress_bar=progress_bar,

tests/test_io_cloud.py

@@ -3,14 +3,14 @@
 
 import pytest
 from litmodels import download_model, upload_model
-from litmodels.io import upload_model_file
+from litmodels.io import upload_model_files
 from torch.nn import Module
 
 
 @pytest.mark.parametrize("name", ["org/model", "model-name", "/too/many/slashes"])
 def test_wrong_model_name(name):
     with pytest.raises(ValueError, match=r".*organization/teamspace/model.*"):
-        upload_model_file(path="path/to/checkpoint", name=name)
+        upload_model_files(path="path/to/checkpoint", name=name)
     with pytest.raises(ValueError, match=r".*organization/teamspace/model.*"):
         download_model(name=name)