Apply suggestions from code review

huggingface · Apr 26, 2022 · 1ffbaff · 1ffbaff
1 parent 1f321d7
commit 1ffbaff
Show file tree

Hide file tree

Showing 5 changed files with 232 additions and 79 deletions.
diff --git a/src/transformers/models/auto/feature_extraction_auto.py b/src/transformers/models/auto/feature_extraction_auto.py
@@ -61,6 +61,7 @@
         ("data2vec-vision", "BeitFeatureExtractor"),
         ("dpt", "DPTFeatureExtractor"),
         ("glpn", "GLPNFeatureExtractor"),
+        ("yolos", "DetrFeatureExtractor"),
     ]
 )
 

diff --git a/src/transformers/models/yolos/configuration_yolos.py b/src/transformers/models/yolos/configuration_yolos.py
@@ -21,7 +21,7 @@
 logger = logging.get_logger(__name__)
 
 YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "google/yolos-s": "https://huggingface.co/yolos-s/resolve/main/config.json",
+    "hustvl/yolos-small": "https://huggingface.co/hustvl/yolos-small/resolve/main/config.json",
     # See all YOLOS models at https://huggingface.co/models?filter=yolos
 }
 
@@ -31,7 +31,7 @@ class YolosConfig(PretrainedConfig):
     This is the configuration class to store the configuration of a [`YolosModel`]. It is used to instantiate a YOLOS
     model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
     defaults will yield a similar configuration to that of the YOLOS
-    [google/vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224) architecture.
+    [hustvl/yolos-base](https://huggingface.co/hustvl/yolos-base) architecture.
 
     Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
     documentation from [`PretrainedConfig`] for more information.
@@ -88,10 +88,10 @@ class YolosConfig(PretrainedConfig):
     ```python
     >>> from transformers import YolosModel, YolosConfig
 
-    >>> # Initializing a YOLOS vit-base-patch16-224 style configuration
+    >>> # Initializing a YOLOS hustvl/yolos-base style configuration
     >>> configuration = YolosConfig()
 
-    >>> # Initializing a model from the vit-base-patch16-224 style configuration
+    >>> # Initializing a model from the hustvl/yolos-base style configuration
     >>> model = YolosModel(configuration)
 
     >>> # Accessing the model configuration

diff --git a/src/transformers/models/yolos/convert_yolos_to_pytorch.py b/src/transformers/models/yolos/convert_yolos_to_pytorch.py
@@ -225,9 +225,18 @@ def convert_yolos_checkpoint(yolos_name, checkpoint_path, pytorch_dump_folder_pa
     feature_extractor.save_pretrained(pytorch_dump_folder_path)
 
     if push_to_hub:
+        model_mapping = {
+            "yolos_ti": "yolos-tiny",
+            "yolos_s_200_pre": "yolos-small",
+            "yolos_s_300_pre": "yolos-small-300",
+            "yolos_s_dWr": "yolos-small-dwr",
+            "yolos_base": "yolos-base",
+        }
+
         print("Pushing to the hub...")
-        feature_extractor.push_to_hub("nielsr/yolos-s")
-        model.push_to_hub("nielsr/yolos-s")
+        model_name = model_mapping[yolos_name]
+        feature_extractor.push_to_hub(model_name, organization="hustvl")
+        model.push_to_hub(model_name, organization="hustvl")
 
 
 if __name__ == "__main__":