cleaning

Author: Borda

src/litmodels/io/__init__.py

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

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_files(
+def upload_model_file(
     name: str,
     path: str,
     progress_bar: bool = True,

tests/test_cloud_io.py

@@ -2,14 +2,15 @@
 from unittest import mock
 
 import pytest
-from litmodels.cloud_io import download_model, upload_model, upload_model_files
+from litmodels import download_model, upload_model
+from litmodels.io import upload_model_file
 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_files(path="path/to/checkpoint", name=name)
+        upload_model_file(path="path/to/checkpoint", name=name)
     with pytest.raises(ValueError, match=r".*organization/teamspace/model.*"):
         download_model(name=name)
 
@@ -25,7 +26,7 @@ def test_wrong_model_name(name):
 def test_upload_model(mocker, tmpdir, model, model_path):
     # mocking the _get_teamspace to return another mock
     ts_mock = mock.MagicMock()
-    mocker.patch("litmodels.cloud_io._get_teamspace", return_value=ts_mock)
+    mocker.patch("litmodels.io.cloud._get_teamspace", return_value=ts_mock)
 
     # The lit-logger function is just a wrapper around the SDK function
     upload_model(