added upload of save_pretrained model

exasol · Nov 29, 2023 · cdcccc7 · cdcccc7
1 parent 2299a8a
commit cdcccc7
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/exasol_transformers_extension/utils/huggingface_hub_bucketfs_model_transfer_sp.py b/exasol_transformers_extension/utils/huggingface_hub_bucketfs_model_transfer_sp.py
@@ -60,7 +60,7 @@ def upload_to_bucketfs(self) -> Path:
         """
         Upload the downloaded models into the BucketFS
         """
-        return self._bucketfs_model_uploader.upload_directory(self._tmpdir_name)
+        return self._bucketfs_model_uploader.upload_directory(self._local_model_save_path / self._model_name)
 
 
 class HuggingFaceHubBucketFSModelTransferSPFactory:

diff --git a/tests/unit_tests/utils/test_huggingface_hub_bucketfs__model_transfer_sp.py b/tests/unit_tests/utils/test_huggingface_hub_bucketfs__model_transfer_sp.py
@@ -83,3 +83,10 @@ def test_download_with_model():
         #todo delete model
 
 
+def test_upload_function_call():
+    test_setup = TestSetup()
+    test_setup.downloader.download_from_huggingface_hub_sp(model_factory=test_setup.model_factory_mock)
+    test_setup.reset_mocks()
+    model_save_path = (test_setup.downloader._local_model_save_path / test_setup.model_name)
+    test_setup.downloader.upload_to_bucketfs()
+    assert mock_cast(test_setup.bucketfs_model_uploader_mock.upload_directory).mock_calls == [call(model_save_path)]