Merge pull request #16 from OpenRL-Lab/huangshiyu

support upload file and directory

Author: huangshiyu13

README.md

@@ -32,6 +32,14 @@ Firstly, you need to login with `huggingface-cli login` (you can create or find
 
 You can also use htool to upload datasets and models to huggingface.
 
+- Upload a file with: `htool upload-file <local_filepath> <organization_or_username/repo_name>:<remote_filepath> -r <model/dataset>`. `-r` means the repo is a model or dataset repo. By default, it is a model repo.
+  - For example: `htool upload-file README.md OpenRL/tizero:README.md`
+  - For example: `htool upload-file README.md OpenRL/DeepFakeFace:README.md -r dataset`
+- Upload a directory with: `htool upload-dir <local_dirpath> <organization_or_username/repo_name>:<remote_dirpath> -r <model/dataset>`. `-r` means the repo is a model or dataset repo. By default, it is a model repo.
+  - For example: `htool upload-dir ./tizero OpenRL/tizero:./`
+  - For example: `htool upload-dir ./tizero/models OpenRL/tizero:./models`
+  - For example: `htool upload-dir ./DeepFakeFace OpenRL/DeepFakeFace:./ -r dataset`
+  - For example: `htool upload-dir ./DeepFakeFace/images OpenRL/DeepFakeFace:./images -r dataset`
 - Upload dataset with: `htool upload-data <local_dataset_dir> <organization_or_username/dataset_name>`
   - For example: `htool upload-data ./daily_dialog OpenRL/daily_dialog`
 - Upload model with: `htool upload-model <local_model_dir> <organization_or_username/model_name>`

huggingface_tool/cli/cli.py

@@ -198,6 +198,53 @@ def upload_data(dataset_dir, dataset_name):
     else:
         uploader.logger.info("Dataset not valid")
 
+@cli.command()
+@click.argument("file_path")
+@click.argument("remote_file_path")
+@click.option(
+    "--repo_type",
+    "-r",
+    type=click.Choice(
+        [
+            "model",
+            "dataset",
+        ]
+    ),
+    default="model",
+    help="repo type",
+)
+def upload_file(file_path, remote_file_path,repo_type):
+    from huggingface_tool.uploaders.file_uploader import FileUploader
+
+    uploader = FileUploader(file_path, remote_file_path,repo_type)
+    if uploader.check():
+        uploader.push()
+    else:
+        uploader.logger.info("File not valid")
+
+@cli.command()
+@click.argument("dir_path")
+@click.argument("remote_dir_path")
+@click.option(
+    "--repo_type",
+    "-r",
+    type=click.Choice(
+        [
+            "model",
+            "dataset",
+        ]
+    ),
+    default="model",
+    help="repo type",
+)
+def upload_dir(dir_path, remote_dir_path,repo_type):
+    from huggingface_tool.uploaders.dir_uploader import DirUploader
+
+    uploader = DirUploader(dir_path, remote_dir_path,repo_type)
+    if uploader.check():
+        uploader.push()
+    else:
+        uploader.logger.info("Directory not valid")
 
 @cli.command()
 @click.argument("model_dir")
@@ -209,4 +256,4 @@ def upload_model(model_dir, model_name):
     if uploader.check():
         uploader.push()
     else:
-        uploader.logger.info("Dataset not valid")
+        uploader.logger.info("Model not valid")

huggingface_tool/uploaders/base_api_uploader.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 2023 The OpenRL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""""""
+from typing import Dict
+from abc import ABC, abstractmethod
+
+from pathlib import Path
+from huggingface_tool.uploaders.base_uploader import BaseUploader
+
+
+class BaseAPIUploader(BaseUploader, ABC):
+    def __init__(self, file_or_dir: str, remote_path: str, repo_type: str):
+        super().__init__(file_or_dir,remote_path)
+        self.info = self.get_info(remote_path)
+        self.repo_type = repo_type
+
+    def get_info(self, remote_path: str) -> Dict[str, str]:
+        split_string = remote_path.split(":")
+        assert len(split_string) == 2
+        repo_id = split_string[0]
+        file_or_directory = split_string[1]
+        return {"repo_id": repo_id, "file_or_directory": file_or_directory}
+
+    def check(self) -> bool:
+        assert Path(
+            self.file_or_dir
+        ).exists(), f"File or directory {self.file_or_dir} does not exist"
+        return True
+

huggingface_tool/uploaders/dir_uploader.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 2023 The OpenRL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""""""
+
+from huggingface_hub import HfApi
+from huggingface_tool.uploaders.base_api_uploader import BaseAPIUploader
+
+
+class DirUploader(BaseAPIUploader):
+    def _push(self) -> bool:
+        try:
+            api = HfApi()
+            api.upload_folder(
+                folder_path=self.file_or_dir,
+                path_in_repo=self.info["file_or_directory"],
+                repo_id=self.info["repo_id"],
+                repo_type=self.repo_type, )
+        except:
+            print("Cannot upload to hub")
+            return False
+        return True
+
+    def _success_message(self):
+        middle_type = "/datasets/" if self.repo_type == "dataset" else "/"
+        print(
+            f"Directory is uploaded to https://huggingface.co{middle_type}{self.info['repo_id']}"
+        )
+
+

huggingface_tool/uploaders/file_uploader.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 2023 The OpenRL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""""""
+
+from huggingface_hub import HfApi
+from huggingface_tool.uploaders.base_api_uploader import BaseAPIUploader
+
+
+class FileUploader(BaseAPIUploader):
+    def _push(self) -> bool:
+        try:
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj=self.file_or_dir,
+                path_in_repo=self.info["file_or_directory"],
+                repo_id=self.info["repo_id"],
+                repo_type=self.repo_type, )
+        except:
+            print("Cannot upload to hub")
+            return False
+        return True
+
+    def _success_message(self):
+        middle_type = "/datasets/" if self.repo_type == "dataset" else "/"
+        print(
+            f"File is uploaded to https://huggingface.co{middle_type}{self.info['repo_id']}"
+        )
+