dump sklearn model to a file for upload (#37)

* dump sklearn model to a file for upload
* linter
* Any

Author: Borda

_requirements/test.txt

@@ -4,3 +4,4 @@ pytest-cov
 pytest-mock
 
 pytorch-lightning >=2.0
+scikit-learn >=1.0

requirements.txt

@@ -2,3 +2,4 @@
 
 lightning-sdk >=0.1.40
 lightning-utilities
+joblib

src/litmodels/io/gateway.py

@@ -1,8 +1,9 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Union
 
+import joblib
 from lightning_utilities import module_available
 
 from litmodels.io.cloud import download_model_files, upload_model_files
@@ -19,7 +20,7 @@
 
 def upload_model(
     name: str,
-    model: Union[str, Path, "Module"],
+    model: Union[str, Path, "Module", Any],
     progress_bar: bool = True,
     cloud_account: Optional[str] = None,
     staging_dir: Optional[str] = None,
@@ -52,7 +53,9 @@ def upload_model(
     elif isinstance(model, Path):
         path = str(model)
     else:
-        raise ValueError(f"Unsupported model type {type(model)}")
+        path = os.path.join(staging_dir, f"{model.__class__.__name__}.pkl")
+        joblib.dump(model, path)
+
     return upload_model_files(
         path=path,
         name=name,

tests/test_io_cloud.py

@@ -4,6 +4,7 @@
 import pytest
 from litmodels import download_model, upload_model
 from litmodels.io import upload_model_files
+from sklearn import svm
 from torch.nn import Module
 
 
@@ -21,6 +22,7 @@ def test_wrong_model_name(name):
         ("path/to/checkpoint", "path/to/checkpoint", False),
         # (BoringModel(), "%s/BoringModel.ckpt"),
         (Module(), f"%s{os.path.sep}Module.pth", True),
+        (svm.SVC(), f"%s{os.path.sep}SVC.pkl", 1),
     ],
 )
 @mock.patch("litmodels.io.cloud.sdk_upload_model")