add save repo and file

add save repo and file

Author: huangshiyu13

Makefile

@@ -1,7 +1,7 @@
 SHELL=/bin/bash
-PROJECT_NAME=openplugin
+PROJECT_NAME=huggingface_tool
 PROJECT_PATH=${PROJECT_NAME}/
-PYTHON_FILES = $(shell find setup.py ${PROJECT_NAME} tests examples -type f -name "*.py")
+PYTHON_FILES = $(shell find setup.py ${PROJECT_NAME} tests -type f -name "*.py")
 
 check_install = python3 -c "import $(1)" || pip3 install $(1) --upgrade
 check_install_extra = python3 -c "import $(1)" || pip3 install $(2) --upgrade

README.md

@@ -15,6 +15,12 @@ Tools for loading, upload, managing huggingface models and datasets
 
 Firstly, you need to login with `huggingface-cli login` (you can create or find your token at [settings](https://huggingface.co/settings/tokens)).
 
+- Download and save a repo with: `htool save-repo <repo_id> <save_dir> -r <model/dataset>`. `-r` means the repo is a model or dataset repo. By default, it is a model repo.
+  - For example: `htool save-repo OpenRL/tizero ./tizero`
+  - For example: `htool save-repo OpenRL/DeepFakeFace ./DeepFakeFace -r dataset`
+- Download and save a file with: `htool save-file <repo_id>:<remote_filepath> <save_dir> -r <model/dataset>`. `-r` means the repo is a model or dataset repo. By default, it is a model repo.
+  - For example: `htool save-file OpenRL/tizero:actor.pt ./tizero`
+  - For example: `htool save-file OpenRL/DeepFakeFace:README.md ./DeepFakeFace -r dataset`
 - Download and save transformer models with: `htool save-model <model_class> <model_name> <save_dir>`
   - For example: `htool save-model AutoModelForCausalLM gpt2 ./gpt2`
 - Download and save tokenizer with: `htool save-tk <tokenizer_name> <save_dir>`

huggingface_tool/cli/cli.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 """"""
-import os
+
 
 import click
 from click.core import Context, Option
@@ -30,9 +30,9 @@ def red(text: str):
 
 
 def print_version(
-        ctx: Context,
-        param: Option,
-        value: bool,
+    ctx: Context,
+    param: Option,
+    value: bool,
 ) -> None:
     if not value or ctx.resilient_parsing:
         return
@@ -42,9 +42,9 @@ def print_version(
 
 
 def print_system_info(
-        ctx: Context,
-        param: Option,
-        value: bool,
+    ctx: Context,
+    param: Option,
+    value: bool,
 ) -> None:
     if not value or ctx.resilient_parsing:
         return
@@ -88,18 +88,70 @@ def cli(ctx):
 @click.argument("save_dir")
 def save_dm(model_name, save_dir):
     from huggingface_tool.savers.diffusion_model_saver import DiffusionModelSaver
+
     saver = DiffusionModelSaver(model_name)
     if saver.load():
         saver.save(save_dir)
     else:
         saver.logger.info("Model not found")
 
 
+@cli.command()
+@click.argument("file_name")
+@click.argument("save_dir")
+@click.option(
+    "--repo_type",
+    "-r",
+    type=click.Choice(
+        [
+            "model",
+            "dataset",
+        ]
+    ),
+    default="model",
+    help="repo type",
+)
+def save_file(file_name, save_dir, repo_type):
+    from huggingface_tool.savers.file_saver import FileSaver
+
+    saver = FileSaver(file_name, repo_type)
+    if saver.load():
+        saver.save(save_dir)
+    else:
+        saver.logger.info("File not found")
+
+
+@cli.command()
+@click.argument("repo_name")
+@click.argument("save_dir")
+@click.option(
+    "--repo_type",
+    "-r",
+    type=click.Choice(
+        [
+            "model",
+            "dataset",
+        ]
+    ),
+    default="model",
+    help="repo type",
+)
+def save_repo(repo_name, save_dir, repo_type):
+    from huggingface_tool.savers.repo_saver import RepoSaver
+
+    saver = RepoSaver(repo_name, repo_type)
+    if saver.load():
+        saver.save(save_dir)
+    else:
+        saver.logger.info("Repo not found")
+
+
 @cli.command()
 @click.argument("tokenizer_name")
 @click.argument("save_dir")
 def save_tk(tokenizer_name, save_dir):
     from huggingface_tool.savers.tokenizer_saver import TokenizerSaver
+
     saver = TokenizerSaver(tokenizer_name)
     if saver.load():
         saver.save(save_dir)
@@ -112,6 +164,7 @@ def save_tk(tokenizer_name, save_dir):
 @click.argument("save_dir")
 def save_data(dataset_name, save_dir):
     from huggingface_tool.savers.dataset_saver import DatasetSaver
+
     saver = DatasetSaver(dataset_name)
     if saver.load():
         saver.save(save_dir)
@@ -125,28 +178,33 @@ def save_data(dataset_name, save_dir):
 @click.argument("save_dir")
 def save_model(model_class, model_name, save_dir):
     from huggingface_tool.savers.model_saver import ModelSaver
+
     saver = ModelSaver(model_class, model_name)
     if saver.load():
         saver.save(save_dir)
     else:
         saver.logger.info("Dataset not found")
 
+
 @cli.command()
 @click.argument("dataset_dir")
 @click.argument("dataset_name")
 def upload_data(dataset_dir, dataset_name):
     from huggingface_tool.uploaders.dataset_uploader import DatasetUploader
+
     uploader = DatasetUploader(dataset_dir, dataset_name)
     if uploader.check():
         uploader.push()
     else:
         uploader.logger.info("Dataset not valid")
 
+
 @cli.command()
 @click.argument("model_dir")
 @click.argument("model_name")
 def upload_model(model_dir, model_name):
     from huggingface_tool.uploaders.model_uploader import ModelUploader
+
     uploader = ModelUploader(model_dir, model_name)
     if uploader.check():
         uploader.push()

huggingface_tool/savers/base_api_saver.py

@@ -0,0 +1,26 @@
+#!/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 abc import ABC
+
+from huggingface_tool.savers.base_saver import BaseSaver
+
+
+class BaseAPISaver(BaseSaver, ABC):
+    def __init__(self, name: str, repo_type: str):
+        super().__init__(name)
+        self.repo_type = repo_type

huggingface_tool/savers/base_model_saver.py

@@ -16,12 +16,13 @@
 
 """"""
 from abc import ABC
+
 from huggingface_tool.savers.base_saver import BaseSaver
 
-class BaseModelSaver(BaseSaver,ABC):
+
+class BaseModelSaver(BaseSaver, ABC):
     def save(self, save_dir: str):
         if self.loaded_object is None:
             self.logger.info("No model loaded, cannot save")
             return
         self.loaded_object.save_pretrained(save_dir)
-

huggingface_tool/savers/base_saver.py

@@ -20,14 +20,13 @@
 
 from huggingface_tool.utils.logger import Logger
 
+
 class BaseSaver(ABC):
-    def __init__(self, name:str):
+    def __init__(self, name: str):
         self.logger = Logger()
         self.name = name
         self.loaded_object = None
 
-
-
     def load(self) -> bool:
         try:
             self.loaded_object = self._load(self.name)
@@ -37,9 +36,9 @@ def load(self) -> bool:
         return True
 
     @abstractmethod
-    def _load(self,name:str):
+    def _load(self, name: str):
         raise NotImplementedError
 
     @abstractmethod
-    def save(self,save_dir: str):
+    def save(self, save_dir: str):
         raise NotImplementedError

huggingface_tool/savers/dataset_saver.py

@@ -19,13 +19,13 @@
 
 from huggingface_tool.savers.base_saver import BaseSaver
 
+
 class DatasetSaver(BaseSaver):
-    def _load(self,name:str):
+    def _load(self, name: str):
         return datasets.load_dataset(name)
 
     def save(self, save_dir: str):
         if self.loaded_object is None:
             self.logger.info("No dataset loaded, cannot save")
             return
         self.loaded_object.save_to_disk(save_dir)
-

huggingface_tool/savers/diffusion_model_saver.py

@@ -15,11 +15,12 @@
 # limitations under the License.
 
 """"""
-from huggingface_tool.savers.base_model_saver import BaseModelSaver
-
 import torch
 from diffusers import StableDiffusionPipeline
 
+from huggingface_tool.savers.base_model_saver import BaseModelSaver
+
+
 class DiffusionModelSaver(BaseModelSaver):
-    def _load(self, name)->bool:
+    def _load(self, name) -> bool:
         return StableDiffusionPipeline.from_pretrained(name, torch_dtype=torch.float16)

huggingface_tool/savers/file_saver.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 pathlib import Path
+
+from huggingface_hub import hf_hub_download
+
+from huggingface_tool.savers.base_api_saver import BaseAPISaver
+
+
+class FileSaver(BaseAPISaver):
+    def _load(self, name: str):
+        split_string = name.split(":")
+        assert len(split_string) == 2
+        repo_id = split_string[0]
+        file = split_string[1]
+        repo_type = self.repo_type if self.repo_type == "dataset" else None
+        return {"repo_id": repo_id, "file": file, "repo_type": repo_type}
+
+    def save(self, name):
+        local_dir = Path(name)
+        local_dir.mkdir(parents=True, exist_ok=True)
+        hf_hub_download(
+            repo_id=self.loaded_object["repo_id"],
+            filename=self.loaded_object["file"],
+            repo_type=self.loaded_object["repo_type"],
+            local_dir=local_dir,
+        )

huggingface_tool/savers/model_saver.py

@@ -22,20 +22,27 @@
 model_class_dict = {
     "AutoModelForSeq2SeqLM": transformers.AutoModelForSeq2SeqLM,
     "AutoModelForCausalLM": transformers.AutoModelForCausalLM,
-    "AutoModelForSequenceClassification": transformers.AutoModelForSequenceClassification,
+    "AutoModelForSequenceClassification": (
+        transformers.AutoModelForSequenceClassification
+    ),
     "AutoModelForQuestionAnswering": transformers.AutoModelForQuestionAnswering,
     "AutoModelForTokenClassification": transformers.AutoModelForTokenClassification,
     "AutoModelForMultipleChoice": transformers.AutoModelForMultipleChoice,
-    "AutoModelForNextSentencePrediction": transformers.AutoModelForNextSentencePrediction,
+    "AutoModelForNextSentencePrediction": (
+        transformers.AutoModelForNextSentencePrediction
+    ),
     "AutoModelForPreTraining": transformers.AutoModelForPreTraining,
     "AutoModelForMaskedLM": transformers.AutoModelForMaskedLM,
-    "AutoModelForTableQuestionAnswering": transformers.AutoModelForTableQuestionAnswering,
+    "AutoModelForTableQuestionAnswering": (
+        transformers.AutoModelForTableQuestionAnswering
+    ),
 }
 
+
 class ModelSaver(BaseModelSaver):
-    def __init__(self, model_class:str, name:str):
+    def __init__(self, model_class: str, name: str):
         super().__init__(name)
         self.model_class = model_class
 
-    def _load(self,name):
-        return model_class_dict[self.model_class].from_pretrained(name)
+    def _load(self, name):
+        return model_class_dict[self.model_class].from_pretrained(name)

huggingface_tool/savers/repo_saver.py

@@ -0,0 +1,38 @@
+#!/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 pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+from huggingface_tool.savers.base_api_saver import BaseAPISaver
+
+
+class RepoSaver(BaseAPISaver):
+    def _load(self, name: str):
+        repo_id = name
+        repo_type = self.repo_type if self.repo_type == "dataset" else None
+        return {"repo_id": repo_id, "repo_type": repo_type}
+
+    def save(self, name):
+        local_dir = Path(name)
+        snapshot_download(
+            repo_id=self.loaded_object["repo_id"],
+            repo_type=self.loaded_object["repo_type"],
+            local_dir=local_dir,
+        )