Update standardize_model_attributes (#1686)

* update files

* update library getter

* fixed typo

* fix library name detection

* fix library name

* update todo

Author: mht-sharma

optimum/exporters/onnx/__main__.py

@@ -211,7 +211,6 @@ def main_export(
     task = TasksManager.map_from_synonym(task)
 
     framework = TasksManager.determine_framework(model_name_or_path, subfolder=subfolder, framework=framework)
-
     library_name = TasksManager.infer_library_from_model(
         model_name_or_path, subfolder=subfolder, library_name=library_name
     )

optimum/exporters/onnx/convert.py

@@ -988,7 +988,7 @@ def onnx_export_from_model(
     """
     library_name = TasksManager._infer_library_from_model(model)
 
-    # TODO: call standardize_model_attributes here once its model_name_or_path argument is optional.
+    TasksManager.standardize_model_attributes(model, library_name)
 
     if hasattr(model.config, "export_model_type"):
         model_type = model.config.export_model_type.replace("_", "-")

optimum/exporters/onnx/model_configs.py

@@ -781,23 +781,6 @@ class TimmDefaultOnnxConfig(ViTOnnxConfig):
     ATOL_FOR_VALIDATION = 1e-3
     DEFAULT_ONNX_OPSET = 12
 
-    def __init__(
-        self,
-        config: "PretrainedConfig",
-        task: str = "feature-extraction",
-        preprocessors: Optional[List[Any]] = None,
-        int_dtype: str = "int64",
-        float_dtype: str = "fp32",
-        legacy: bool = False,
-    ):
-        super().__init__(config, task, preprocessors, int_dtype, float_dtype, legacy)
-
-        pretrained_cfg = self._config
-        if hasattr(self._config, "pretrained_cfg"):
-            pretrained_cfg = self._config.pretrained_cfg
-
-        self._normalized_config = self.NORMALIZED_CONFIG_CLASS(pretrained_cfg)
-
     def rename_ambiguous_inputs(self, inputs):
         #  The input name in the model signature is `x, hence the export input name is updated.
         model_inputs = {}

optimum/exporters/tasks.py

@@ -17,7 +17,6 @@
 import importlib
 import inspect
 import itertools
-import json
 import os
 from functools import partial
 from pathlib import Path
@@ -1523,10 +1522,12 @@ def _infer_task_from_model_name_or_path(
                     "Cannot infer the task from a model repo with a subfolder yet, please specify the task manually."
                 )
             model_info = huggingface_hub.model_info(model_name_or_path, revision=revision)
-            if getattr(model_info, "library_name", None) == "diffusers":
+            library_name = TasksManager.infer_library_from_model(model_name_or_path, subfolder, revision)
+
+            if library_name == "diffusers":
                 class_name = model_info.config["diffusers"]["class_name"]
                 inferred_task_name = "stable-diffusion-xl" if "StableDiffusionXL" in class_name else "stable-diffusion"
-            elif getattr(model_info, "library_name", None) == "timm":
+            elif library_name == "timm":
                 inferred_task_name = "image-classification"
             else:
                 pipeline_tag = getattr(model_info, "pipeline_tag", None)
@@ -1544,13 +1545,9 @@ def _infer_task_from_model_name_or_path(
                         # transformersInfo does not always have a pipeline_tag attribute
                         class_name_prefix = ""
                         if is_torch_available():
-                            tasks_to_automodels = TasksManager._LIBRARY_TO_TASKS_TO_MODEL_LOADER_MAP[
-                                model_info.library_name
-                            ]
+                            tasks_to_automodels = TasksManager._LIBRARY_TO_TASKS_TO_MODEL_LOADER_MAP[library_name]
                         else:
-                            tasks_to_automodels = TasksManager._LIBRARY_TO_TF_TASKS_TO_MODEL_LOADER_MAP[
-                                model_info.library_name
-                            ]
+                            tasks_to_automodels = TasksManager._LIBRARY_TO_TF_TASKS_TO_MODEL_LOADER_MAP[library_name]
                             class_name_prefix = "TF"
 
                         auto_model_class_name = transformers_info["auto_model"]
@@ -1603,8 +1600,17 @@ def infer_task_from_model(
         return task
 
     @staticmethod
-    def _infer_library_from_model(model: Union["PreTrainedModel", "TFPreTrainedModel"]):
-        if hasattr(model.config, "pretrained_cfg") or hasattr(model.config, "architecture"):
+    def _infer_library_from_model(
+        model: Union["PreTrainedModel", "TFPreTrainedModel"], library_name: Optional[str] = None
+    ):
+        if library_name is not None:
+            return library_name
+
+        if (
+            hasattr(model, "pretrained_cfg")
+            or hasattr(model.config, "pretrained_cfg")
+            or hasattr(model.config, "architecture")
+        ):
             library_name = "timm"
         elif hasattr(model.config, "_diffusers_version") or getattr(model, "config_name", "") == "model_index.json":
             library_name = "diffusers"
@@ -1645,44 +1651,33 @@ def infer_library_from_model(
         if library_name is not None:
             return library_name
 
-        full_model_path = Path(model_name_or_path) / subfolder
-
-        if not full_model_path.is_dir():
-            model_info = huggingface_hub.model_info(model_name_or_path, revision=revision)
-            library_name = getattr(model_info, "library_name", None)
-
-            # sentence-transformers package name is sentence_transformers
-            if library_name is not None:
-                library_name = library_name.replace("-", "_")
+        all_files, _ = TasksManager.get_model_files(model_name_or_path, subfolder, cache_dir)
 
-        if library_name is None:
-            all_files, _ = TasksManager.get_model_files(model_name_or_path, subfolder, cache_dir)
+        if "model_index.json" in all_files:
+            library_name = "diffusers"
+        elif CONFIG_NAME in all_files:
+            # We do not use PretrainedConfig.from_pretrained which has unwanted warnings about model type.
+            kwargs = {
+                "subfolder": subfolder,
+                "revision": revision,
+                "cache_dir": cache_dir,
+            }
+            config_dict, kwargs = PretrainedConfig.get_config_dict(model_name_or_path, **kwargs)
+            model_config = PretrainedConfig.from_dict(config_dict, **kwargs)
 
-            if "model_index.json" in all_files:
+            if hasattr(model_config, "pretrained_cfg") or hasattr(model_config, "architecture"):
+                library_name = "timm"
+            elif hasattr(model_config, "_diffusers_version"):
                 library_name = "diffusers"
-            elif CONFIG_NAME in all_files:
-                # We do not use PretrainedConfig.from_pretrained which has unwanted warnings about model type.
-                kwargs = {
-                    "subfolder": subfolder,
-                    "revision": revision,
-                    "cache_dir": cache_dir,
-                }
-                config_dict, kwargs = PretrainedConfig.get_config_dict(model_name_or_path, **kwargs)
-                model_config = PretrainedConfig.from_dict(config_dict, **kwargs)
-
-                if hasattr(model_config, "pretrained_cfg") or hasattr(model_config, "architecture"):
-                    library_name = "timm"
-                elif hasattr(model_config, "_diffusers_version"):
-                    library_name = "diffusers"
-                else:
-                    library_name = "transformers"
-            elif (
-                any(file_path.startswith("sentence_") for file_path in all_files)
-                or "config_sentence_transformers.json" in all_files
-            ):
-                library_name = "sentence_transformers"
             else:
                 library_name = "transformers"
+        elif (
+            any(file_path.startswith("sentence_") for file_path in all_files)
+            or "config_sentence_transformers.json" in all_files
+        ):
+            library_name = "sentence_transformers"
+        else:
+            library_name = "transformers"
 
         if library_name is None:
             raise ValueError(
@@ -1694,11 +1689,7 @@ def infer_library_from_model(
     @classmethod
     def standardize_model_attributes(
         cls,
-        model_name_or_path: Union[str, Path],
         model: Union["PreTrainedModel", "TFPreTrainedModel"],
-        subfolder: str = "",
-        revision: Optional[str] = None,
-        cache_dir: str = huggingface_hub.constants.HUGGINGFACE_HUB_CACHE,
         library_name: Optional[str] = None,
     ):
         """
@@ -1721,40 +1712,20 @@ def standardize_model_attributes(
             library_name (`Optional[str]`, *optional*)::
                 The library name of the model. Can be any of "transformers", "timm", "diffusers", "sentence_transformers".
         """
-        # TODO: make model_name_or_path an optional argument here.
-
-        library_name = TasksManager.infer_library_from_model(
-            model_name_or_path, subfolder, revision, cache_dir, library_name
-        )
-
-        full_model_path = Path(model_name_or_path) / subfolder
-        is_local = full_model_path.is_dir()
+        library_name = TasksManager._infer_library_from_model(model, library_name)
 
         if library_name == "diffusers":
             model.config.export_model_type = "stable-diffusion"
         elif library_name == "timm":
             # Retrieve model config
-            config_path = full_model_path / "config.json"
-
-            if not is_local:
-                config_path = huggingface_hub.hf_hub_download(
-                    model_name_or_path, "config.json", subfolder=subfolder, revision=revision
-                )
-
-            model_config = PretrainedConfig.from_json_file(config_path)
-
-            if hasattr(model_config, "pretrained_cfg"):
-                model_config.pretrained_cfg = PretrainedConfig.from_dict(model_config.pretrained_cfg)
+            model_config = PretrainedConfig.from_dict(model.pretrained_cfg)
 
             # Set config as in transformers
             setattr(model, "config", model_config)
 
-            # Update model_type for model
-            with open(config_path) as fp:
-                model_type = json.load(fp)["architecture"]
-
             # `model_type` is a class attribute in Transformers, let's avoid modifying it.
-            model.config.export_model_type = model_type
+            model.config.export_model_type = model.pretrained_cfg["architecture"]
+
         elif library_name == "sentence_transformers":
             if "Transformer" in model[0].__class__.__name__:
                 model.config = model[0].auto_model.config
@@ -1903,9 +1874,7 @@ def get_model_from_task(
                     kwargs["from_pt"] = True
                     model = model_class.from_pretrained(model_name_or_path, **kwargs)
 
-        TasksManager.standardize_model_attributes(
-            model_name_or_path, model, subfolder, revision, cache_dir, library_name
-        )
+        TasksManager.standardize_model_attributes(model, library_name)
 
         return model
 

tests/exporters/onnx/test_onnx_export.py

@@ -185,7 +185,7 @@ def _onnx_export(
         if library_name == "timm":
             model_class = TasksManager.get_model_class_for_task(task, library=library_name)
             model = model_class(f"hf_hub:{model_name}", pretrained=True, exportable=True)
-            TasksManager.standardize_model_attributes(model_name, model, library_name=library_name)
+            TasksManager.standardize_model_attributes(model, library_name=library_name)
         else:
             config = AutoConfig.from_pretrained(model_name)
             model_class = TasksManager.get_model_class_for_task(task, model_type=config.model_type.replace("_", "-"))
@@ -611,7 +611,7 @@ def _onnx_export(
         if library_name == "timm":
             model_class = TasksManager.get_model_class_for_task(task, library=library_name)
             model = model_class(f"hf_hub:{model_name}", pretrained=True, exportable=True)
-            TasksManager.standardize_model_attributes(model_name, model, library_name=library_name)
+            TasksManager.standardize_model_attributes(model, library_name=library_name)
         else:
             config = AutoConfig.from_pretrained(model_name)
             model_class = TasksManager.get_model_class_for_task(task, model_type=config.model_type.replace("_", "-"))