From cdcccc77dbc0181869858dd7acb1add901f2099e Mon Sep 17 00:00:00 2001 From: MarleneKress79789 Date: Wed, 29 Nov 2023 13:43:11 +0100 Subject: [PATCH] added upload of save_pretrained model --- .../utils/huggingface_hub_bucketfs_model_transfer_sp.py | 2 +- .../test_huggingface_hub_bucketfs__model_transfer_sp.py | 7 +++++++ 2 files changed, 8 insertions(+), 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 index 49d8e31a..f9e1a73b 100644 --- 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 index 1cae7766..fd3dc894 100644 --- 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)] \ No newline at end of file