We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
"disable_shape_inference" in ORTOptimizer does not work
### System Info

```shell
optimum==1.6.1
```


### Who can help?

@echarlaix @philschmid 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Using following opt. config:
```
opt_conf = OptimizationConfig(optimization_level=0)
opt_conf.fp16 = True
opt_conf.enable_transformers_specific_optimizations = False
opt_conf.disable_embed_layer_norm_fusion = False
opt_conf.disable_shape_inference = True
```
The error keeps appearing (GPT-J model): `ZeroDivisionError: integer division or modulo by zero`

Error trace contains this:
```
File ~/miniconda3/envs/pytorch/lib/python3.10/site-packages/onnxruntime/transformers/models/gpt2/../../onnx_model.py:594, in OnnxModel.convert_float_to_float16(self, use_symbolic_shape_infer, **kwargs)
    591 if use_symbolic_shape_infer:
    592     # Use symbolic shape inference since custom operators (like Gelu, SkipLayerNormalization etc) are not recognized by onnx shape inference.
    593     shape_infer_helper = SymbolicShapeInferenceHelper(model)
--> 594     model = shape_infer_helper.infer_shapes(model, auto_merge=True, guess_output_rank=False)
    596 parameters = {"disable_shape_infer": use_symbolic_shape_infer}
    597 parameters.update(
    598     {
    599         key: kwargs[key]
   (...)
    609     }
    610 )
```

And if you check the code in optimum/onnxruntime/optimization.py on line number 159 you will find that
`optimizer.convert_float_to_float16(keep_io_types=True)` doesn't get a parameter that disables symbolic shape inference.

However, the signature of this method is:
`def convert_float_to_float16(self, use_symbolic_shape_infer=True, **kwargs):`

### Expected behavior

I would like the errors with shape inference to be better explained and to be able to turn it off for real.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `optimum/onnxruntime/optimization.py`
Content:
```
1 #  Copyright 2021 The HuggingFace Team. All rights reserved.
2 #
3 #  Licensed under the Apache License, Version 2.0 (the "License");
4 #  you may not use this file except in compliance with the License.
5 #  You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 #  Unless required by applicable law or agreed to in writing, software
10 #  distributed under the License is distributed on an "AS IS" BASIS,
11 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 #  See the License for the specific language governing permissions and
13 #  limitations under the License.
14 """Main class for performing graph optimization with ONNX Runtime."""
15 
16 import logging
17 import os
18 from pathlib import Path
19 from typing import TYPE_CHECKING, Dict, List, Optional, Union
20 
21 from transformers.models.auto.configuration_auto import AutoConfig
22 
23 from onnx import load_model
24 from onnxruntime.transformers.onnx_model_bert import BertOnnxModel
25 from onnxruntime.transformers.optimizer import optimize_model
26 
27 from ..utils import CONFIG_NAME, NormalizedConfigManager
28 from ..utils.save_utils import maybe_save_preprocessors
29 from .configuration import OptimizationConfig, ORTConfig
30 from .modeling_ort import ORTModel
31 from .modeling_seq2seq import ORTModelForSeq2SeqLM
32 from .utils import ONNX_WEIGHTS_NAME, ORTConfigManager
33 
34 
35 if TYPE_CHECKING:
36     from transformers import PretrainedConfig
37 
38 
39 LOGGER = logging.getLogger(__name__)
40 
41 
42 class ORTOptimizer:
43     """
44     Handles the ONNX Runtime optimization process for models shared on huggingface.co/models.
45     """
46 
47     def __init__(self, onnx_model_path: List[os.PathLike], config: "PretrainedConfig"):
48         """
49         Args:
50             onnx_model_path (`List[os.PathLike]`):
51                 The paths of the onnx models to optimize.
52             config ([`~transformers.PretrainedConfig`]):
53                 An instance of the configuration associated to the model to optimize.
54         """
55         super().__init__()
56         self.onnx_model_path = onnx_model_path
57         self.config = config
58         self.model_type = self.config.model_type
59         self.normalized_config = NormalizedConfigManager.get_normalized_config_class(self.model_type)(self.config)
60 
61     @classmethod
62     def from_pretrained(
63         cls, model_or_path: Union[str, os.PathLike, ORTModel], file_names: Optional[List[str]] = None
64     ) -> "ORTOptimizer":
65         """
66         Args:
67             model_or_path (`Union[str, os.PathLike, ORTModel]`):
68                 The path to a local directory hosting the model to optimize or an instance of an `ORTModel` to quantize.
69                 Can be either:
70                     - A path to a local *directory* containing the model to optimize.
71                     - An instance of [`~optimum.onnxruntime.ORTModel`].
72             file_names(`Optional[List[str]]`, *optional*):
73                 The list of file names of the models to optimize.
74         """
75         onnx_model_path = []
76         config = None
77         if isinstance(model_or_path, ORTModel):
78             if isinstance(model_or_path, ORTModelForSeq2SeqLM):
79                 onnx_model_path += [
80                     model_or_path.encoder_model_path,
81                     model_or_path.decoder_model_path,
82                 ]
83                 # Add the decoder with past key/values if present
84                 if model_or_path.use_cache:
85                     onnx_model_path.append(model_or_path.decoder_with_past_model_path)
86             else:
87                 onnx_model_path.append(model_or_path.model_path)
88             config = model_or_path.config
89         elif os.path.isdir(model_or_path):
90             file_names = [ONNX_WEIGHTS_NAME] if file_names is None else file_names
91             model_or_path = Path(model_or_path)
92             if CONFIG_NAME not in os.listdir(model_or_path):
93                 raise ValueError(f"The local directory does not contain the configuration file {CONFIG_NAME}.")
94             config = AutoConfig.from_pretrained(model_or_path)
95             for file_name in file_names:
96                 onnx_model_path.append(model_or_path.joinpath(file_name))
97         else:
98             raise ValueError(f"Unable to load the model from {model_or_path}.")
99         return cls(onnx_model_path, config=config)
100 
101     def optimize(
102         self,
103         optimization_config: OptimizationConfig,
104         save_dir: Union[str, os.PathLike],
105         file_suffix: Optional[str] = "optimized",
106         use_external_data_format: bool = False,
107         one_external_file: bool = True,
108     ):
109         """
110         Optimizes a model given the optimization specifications defined in `optimization_config`.
111 
112         Args:
113             optimization_config ([`~optimum.onnxruntime.OptimizationConfig`]):
114                 The configuration containing the parameters related to optimization.
115             save_dir (`Union[str, os.PathLike]`):
116                 The path used to save the optimized model.
117             file_suffix (`str`, *optional*, defaults to `"optimized"`):
118                 The file suffix used to save the optimized model.
119             use_external_data_format (`bool`, *optional*, defaults to `False`):
120                 Whether to use external data format to store model of size >= 2Gb.
121             one_external_file (`bool`, defaults to `True`):
122                 When `use_external_data_format=True`, whether to save all tensors to one external file.
123                 If false, save each tensor to a file named with the tensor name.
124 
125         """
126         save_dir = Path(save_dir)
127         save_dir.mkdir(parents=True, exist_ok=True)
128         ORTConfigManager.check_optimization_supported_model(self.model_type)
129 
130         self.config.save_pretrained(save_dir)
131         maybe_save_preprocessors(self.onnx_model_path[0].parent, save_dir)
132 
133         # Create and save the configuration summarizing all the parameters related to optimization
134         ort_config = ORTConfig(
135             optimization=optimization_config,
136             use_external_data_format=use_external_data_format,
137             one_external_file=one_external_file,
138         )
139 
140         model_type = ORTConfigManager.get_model_ort_type(self.config.model_type)
141         optimization_options = optimization_config.create_fusion_options(model_type)
142 
143         LOGGER.info("Optimizing model...")
144 
145         for model_path in self.onnx_model_path:
146             optimizer = optimize_model(
147                 model_path.as_posix(),
148                 model_type,
149                 self.normalized_config.num_attention_heads,
150                 self.normalized_config.hidden_size,
151                 opt_level=optimization_config.optimization_level,
152                 optimization_options=optimization_options,
153                 use_gpu=optimization_config.optimize_for_gpu,
154                 only_onnxruntime=not optimization_config.enable_transformers_specific_optimizations,
155             )
156 
157             if optimization_config.fp16:
158                 # keep_io_types to keep inputs/outputs as float32
159                 optimizer.convert_float_to_float16(keep_io_types=True)
160 
161             suffix = f"_{file_suffix}" if file_suffix else ""
162             output_path = save_dir.joinpath(f"{model_path.stem}{suffix}").with_suffix(model_path.suffix)
163             optimizer.save_model_to_file(output_path.as_posix(), use_external_data_format, one_external_file)
164 
165         # Save the model configuration
166         self.config.save_pretrained(save_dir)
167         ort_config.save_pretrained(save_dir)
168 
169         LOGGER.info(
170             f"Optimized model saved at: {save_dir} (external data format: "
171             f"{use_external_data_format}; saved all tensor to one file: "
172             f"{one_external_file})"
173         )
174 
175         return Path(save_dir)
176 
177     @staticmethod
178     def get_fused_operators(onnx_model_path: Union[str, os.PathLike]) -> Dict[str, int]:
179         """
180         Computes the dictionary mapping the name of the fused operators to their number of apparition in the model.
181 
182         Args:
183             onnx_model_path (`Union[str, os.PathLike]`):
184                 Path of the ONNX model.
185 
186         Returns:
187             The dictionary mapping the name of the fused operators to their number of apparition in the model.
188         """
189         onnx_optimized_model = BertOnnxModel(load_model(onnx_model_path))
190         fused_operator = onnx_optimized_model.get_fused_operator_statistics()
191         LOGGER.info(
192             f"The following operators were fused : { ', '.join([k for k,v in fused_operator.items() if v > 0])}"
193         )
194         return {k: v for k, v in fused_operator.items() if v > 0}
195 
196     @staticmethod
197     def get_nodes_number_difference(
198         onnx_model_path: Union[str, os.PathLike], onnx_optimized_model_path: Union[str, os.PathLike]
199     ) -> int:
200         """
201         Compute the difference in the number of nodes between the original and the optimized model.
202 
203         Args:
204             onnx_model_path (`Union[str, os.PathLike]`):
205                 Path of the ONNX model.
206             onnx_optimized_model_path (`Union[str, os.PathLike]`):
207                 Path of the optimized ONNX model.
208 
209         Returns:
210             The difference in the number of nodes between the original and the optimized model.
211         """
212         onnx_model = BertOnnxModel(load_model(onnx_model_path))
213         onnx_optimized_model = BertOnnxModel(load_model(onnx_optimized_model_path))
214 
215         # Information in the number of nodes decrease resulting from optimization
216         nodes_number_onnx_model = len(onnx_model.nodes())
217         nodes_number_onnx_optimized_model = len(onnx_optimized_model.nodes())
218         difference_nodes_number = nodes_number_onnx_model - nodes_number_onnx_optimized_model
219         LOGGER.info(
220             f"There are {nodes_number_onnx_model} nodes before optimization and {nodes_number_onnx_optimized_model}"
221             f"nodes after. The number of nodes removed is {difference_nodes_number}"
222         )
223         return difference_nodes_number
224 
225     @staticmethod
226     def get_operators_difference(
227         onnx_model_path: Union[str, os.PathLike], onnx_optimized_model_path: Union[str, os.PathLike]
228     ) -> Dict[str, int]:
229         """
230         Compute the dictionary mapping the operators name to the difference in the number of corresponding nodes between
231         the original and the optimized model.
232 
233         Args:
234             onnx_model_path (`Union[str, os.PathLike]`):
235                 Path of the ONNX model.
236             onnx_optimized_model_path (`Union[str, os.PathLike]`):
237                 Path of the optimized ONNX model.
238 
239         Returns:
240             The dictionary mapping the operators name to the difference in the number of corresponding nodes between the
241             original and the optimized model.
242         """
243         onnx_model = BertOnnxModel(load_model(onnx_model_path))
244         onnx_optimized_model = BertOnnxModel(load_model(onnx_optimized_model_path))
245 
246         def nodes_difference_given_type(op_type):
247             onnx_model_nodes_with_op_type = len(onnx_model.get_nodes_by_op_type(op_type))
248             onnx_optimized_model_nodes_with_op_type = len(onnx_optimized_model.get_nodes_by_op_type(op_type))
249             return onnx_model_nodes_with_op_type - onnx_optimized_model_nodes_with_op_type
250 
251         # Compute operators difference between the original and the optimized models
252         op_types = set()
253         for model in [onnx_model, onnx_optimized_model]:
254             for node in model.nodes():
255                 op_types.add(node.op_type)
256 
257         operators_difference = dict(map(lambda op_type: (op_type, nodes_difference_given_type(op_type)), op_types))
258         return {k: v for k, v in operators_difference.items() if v != 0}
259
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/optimum/onnxruntime/optimization.py b/optimum/onnxruntime/optimization.py
--- a/optimum/onnxruntime/optimization.py
+++ b/optimum/onnxruntime/optimization.py
@@ -143,20 +143,30 @@
         LOGGER.info("Optimizing model...")
 
         for model_path in self.onnx_model_path:
-            optimizer = optimize_model(
-                model_path.as_posix(),
-                model_type,
-                self.normalized_config.num_attention_heads,
-                self.normalized_config.hidden_size,
-                opt_level=optimization_config.optimization_level,
-                optimization_options=optimization_options,
-                use_gpu=optimization_config.optimize_for_gpu,
-                only_onnxruntime=not optimization_config.enable_transformers_specific_optimizations,
-            )
-
-            if optimization_config.fp16:
-                # keep_io_types to keep inputs/outputs as float32
-                optimizer.convert_float_to_float16(keep_io_types=True)
+            try:
+                optimizer = optimize_model(
+                    model_path.as_posix(),
+                    model_type,
+                    self.normalized_config.num_attention_heads,
+                    self.normalized_config.hidden_size,
+                    opt_level=optimization_config.optimization_level,
+                    optimization_options=optimization_options,
+                    use_gpu=optimization_config.optimize_for_gpu,
+                    only_onnxruntime=not optimization_config.enable_transformers_specific_optimizations,
+                )
+
+                if optimization_config.fp16:
+                    # keep_io_types to keep inputs/outputs as float32
+                    optimizer.convert_float_to_float16(
+                        use_symbolic_shape_infer=not optimization_config.disable_shape_inference, keep_io_types=True
+                    )
+            except Exception as e:
+                if "Incomplete symbolic shape inference" in str(e):
+                    err = RuntimeError(
+                        f"{str(e)}. Try to set `disable_shape_inference=True` in your optimization configuration."
+                    )
+                    raise err from e
+                raise
 
             suffix = f"_{file_suffix}" if file_suffix else ""
             output_path = save_dir.joinpath(f"{model_path.stem}{suffix}").with_suffix(model_path.suffix)

{"golden_diff": "diff --git a/optimum/onnxruntime/optimization.py b/optimum/onnxruntime/optimization.py\n--- a/optimum/onnxruntime/optimization.py\n+++ b/optimum/onnxruntime/optimization.py\n@@ -143,20 +143,30 @@\n         LOGGER.info(\"Optimizing model...\")\n \n         for model_path in self.onnx_model_path:\n-            optimizer = optimize_model(\n-                model_path.as_posix(),\n-                model_type,\n-                self.normalized_config.num_attention_heads,\n-                self.normalized_config.hidden_size,\n-                opt_level=optimization_config.optimization_level,\n-                optimization_options=optimization_options,\n-                use_gpu=optimization_config.optimize_for_gpu,\n-                only_onnxruntime=not optimization_config.enable_transformers_specific_optimizations,\n-            )\n-\n-            if optimization_config.fp16:\n-                # keep_io_types to keep inputs/outputs as float32\n-                optimizer.convert_float_to_float16(keep_io_types=True)\n+            try:\n+                optimizer = optimize_model(\n+                    model_path.as_posix(),\n+                    model_type,\n+                    self.normalized_config.num_attention_heads,\n+                    self.normalized_config.hidden_size,\n+                    opt_level=optimization_config.optimization_level,\n+                    optimization_options=optimization_options,\n+                    use_gpu=optimization_config.optimize_for_gpu,\n+                    only_onnxruntime=not optimization_config.enable_transformers_specific_optimizations,\n+                )\n+\n+                if optimization_config.fp16:\n+                    # keep_io_types to keep inputs/outputs as float32\n+                    optimizer.convert_float_to_float16(\n+                        use_symbolic_shape_infer=not optimization_config.disable_shape_inference, keep_io_types=True\n+                    )\n+            except Exception as e:\n+                if \"Incomplete symbolic shape inference\" in str(e):\n+                    err = RuntimeError(\n+                        f\"{str(e)}. Try to set `disable_shape_inference=True` in your optimization configuration.\"\n+                    )\n+                    raise err from e\n+                raise\n \n             suffix = f\"_{file_suffix}\" if file_suffix else \"\"\n             output_path = save_dir.joinpath(f\"{model_path.stem}{suffix}\").with_suffix(model_path.suffix)\n", "issue": "\"disable_shape_inference\" in ORTOptimizer does not work\n### System Info\r\n\r\n```shell\r\noptimum==1.6.1\r\n```\r\n\r\n\r\n### Who can help?\r\n\r\n@echarlaix @philschmid \r\n\r\n### Information\r\n\r\n- [X] The official example scripts\r\n- [ ] My own modified scripts\r\n\r\n### Tasks\r\n\r\n- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)\r\n- [ ] My own task or dataset (give details below)\r\n\r\n### Reproduction\r\n\r\nUsing following opt. config:\r\n```\r\nopt_conf = OptimizationConfig(optimization_level=0)\r\nopt_conf.fp16 = True\r\nopt_conf.enable_transformers_specific_optimizations = False\r\nopt_conf.disable_embed_layer_norm_fusion = False\r\nopt_conf.disable_shape_inference = True\r\n```\r\nThe error keeps appearing (GPT-J model): `ZeroDivisionError: integer division or modulo by zero`\r\n\r\nError trace contains this:\r\n```\r\nFile ~/miniconda3/envs/pytorch/lib/python3.10/site-packages/onnxruntime/transformers/models/gpt2/../../onnx_model.py:594, in OnnxModel.convert_float_to_float16(self, use_symbolic_shape_infer, **kwargs)\r\n    591 if use_symbolic_shape_infer:\r\n    592     # Use symbolic shape inference since custom operators (like Gelu, SkipLayerNormalization etc) are not recognized by onnx shape inference.\r\n    593     shape_infer_helper = SymbolicShapeInferenceHelper(model)\r\n--> 594     model = shape_infer_helper.infer_shapes(model, auto_merge=True, guess_output_rank=False)\r\n    596 parameters = {\"disable_shape_infer\": use_symbolic_shape_infer}\r\n    597 parameters.update(\r\n    598     {\r\n    599         key: kwargs[key]\r\n   (...)\r\n    609     }\r\n    610 )\r\n```\r\n\r\nAnd if you check the code in optimum/onnxruntime/optimization.py on line number 159 you will find that\r\n`optimizer.convert_float_to_float16(keep_io_types=True)` doesn't get a parameter that disables symbolic shape inference.\r\n\r\nHowever, the signature of this method is:\r\n`def convert_float_to_float16(self, use_symbolic_shape_infer=True, **kwargs):`\r\n\r\n### Expected behavior\r\n\r\nI would like the errors with shape inference to be better explained and to be able to turn it off for real.\n", "before_files": [{"content": "#  Copyright 2021 The HuggingFace Team. All rights reserved.\n#\n#  Licensed under the Apache License, Version 2.0 (the \"License\");\n#  you may not use this file except in compliance with the License.\n#  You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n#  Unless required by applicable law or agreed to in writing, software\n#  distributed under the License is distributed on an \"AS IS\" BASIS,\n#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n#  See the License for the specific language governing permissions and\n#  limitations under the License.\n\"\"\"Main class for performing graph optimization with ONNX Runtime.\"\"\"\n\nimport logging\nimport os\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Dict, List, Optional, Union\n\nfrom transformers.models.auto.configuration_auto import AutoConfig\n\nfrom onnx import load_model\nfrom onnxruntime.transformers.onnx_model_bert import BertOnnxModel\nfrom onnxruntime.transformers.optimizer import optimize_model\n\nfrom ..utils import CONFIG_NAME, NormalizedConfigManager\nfrom ..utils.save_utils import maybe_save_preprocessors\nfrom .configuration import OptimizationConfig, ORTConfig\nfrom .modeling_ort import ORTModel\nfrom .modeling_seq2seq import ORTModelForSeq2SeqLM\nfrom .utils import ONNX_WEIGHTS_NAME, ORTConfigManager\n\n\nif TYPE_CHECKING:\n    from transformers import PretrainedConfig\n\n\nLOGGER = logging.getLogger(__name__)\n\n\nclass ORTOptimizer:\n    \"\"\"\n    Handles the ONNX Runtime optimization process for models shared on huggingface.co/models.\n    \"\"\"\n\n    def __init__(self, onnx_model_path: List[os.PathLike], config: \"PretrainedConfig\"):\n        \"\"\"\n        Args:\n            onnx_model_path (`List[os.PathLike]`):\n                The paths of the onnx models to optimize.\n            config ([`~transformers.PretrainedConfig`]):\n                An instance of the configuration associated to the model to optimize.\n        \"\"\"\n        super().__init__()\n        self.onnx_model_path = onnx_model_path\n        self.config = config\n        self.model_type = self.config.model_type\n        self.normalized_config = NormalizedConfigManager.get_normalized_config_class(self.model_type)(self.config)\n\n    @classmethod\n    def from_pretrained(\n        cls, model_or_path: Union[str, os.PathLike, ORTModel], file_names: Optional[List[str]] = None\n    ) -> \"ORTOptimizer\":\n        \"\"\"\n        Args:\n            model_or_path (`Union[str, os.PathLike, ORTModel]`):\n                The path to a local directory hosting the model to optimize or an instance of an `ORTModel` to quantize.\n                Can be either:\n                    - A path to a local *directory* containing the model to optimize.\n                    - An instance of [`~optimum.onnxruntime.ORTModel`].\n            file_names(`Optional[List[str]]`, *optional*):\n                The list of file names of the models to optimize.\n        \"\"\"\n        onnx_model_path = []\n        config = None\n        if isinstance(model_or_path, ORTModel):\n            if isinstance(model_or_path, ORTModelForSeq2SeqLM):\n                onnx_model_path += [\n                    model_or_path.encoder_model_path,\n                    model_or_path.decoder_model_path,\n                ]\n                # Add the decoder with past key/values if present\n                if model_or_path.use_cache:\n                    onnx_model_path.append(model_or_path.decoder_with_past_model_path)\n            else:\n                onnx_model_path.append(model_or_path.model_path)\n            config = model_or_path.config\n        elif os.path.isdir(model_or_path):\n            file_names = [ONNX_WEIGHTS_NAME] if file_names is None else file_names\n            model_or_path = Path(model_or_path)\n            if CONFIG_NAME not in os.listdir(model_or_path):\n                raise ValueError(f\"The local directory does not contain the configuration file {CONFIG_NAME}.\")\n            config = AutoConfig.from_pretrained(model_or_path)\n            for file_name in file_names:\n                onnx_model_path.append(model_or_path.joinpath(file_name))\n        else:\n            raise ValueError(f\"Unable to load the model from {model_or_path}.\")\n        return cls(onnx_model_path, config=config)\n\n    def optimize(\n        self,\n        optimization_config: OptimizationConfig,\n        save_dir: Union[str, os.PathLike],\n        file_suffix: Optional[str] = \"optimized\",\n        use_external_data_format: bool = False,\n        one_external_file: bool = True,\n    ):\n        \"\"\"\n        Optimizes a model given the optimization specifications defined in `optimization_config`.\n\n        Args:\n            optimization_config ([`~optimum.onnxruntime.OptimizationConfig`]):\n                The configuration containing the parameters related to optimization.\n            save_dir (`Union[str, os.PathLike]`):\n                The path used to save the optimized model.\n            file_suffix (`str`, *optional*, defaults to `\"optimized\"`):\n                The file suffix used to save the optimized model.\n            use_external_data_format (`bool`, *optional*, defaults to `False`):\n                Whether to use external data format to store model of size >= 2Gb.\n            one_external_file (`bool`, defaults to `True`):\n                When `use_external_data_format=True`, whether to save all tensors to one external file.\n                If false, save each tensor to a file named with the tensor name.\n\n        \"\"\"\n        save_dir = Path(save_dir)\n        save_dir.mkdir(parents=True, exist_ok=True)\n        ORTConfigManager.check_optimization_supported_model(self.model_type)\n\n        self.config.save_pretrained(save_dir)\n        maybe_save_preprocessors(self.onnx_model_path[0].parent, save_dir)\n\n        # Create and save the configuration summarizing all the parameters related to optimization\n        ort_config = ORTConfig(\n            optimization=optimization_config,\n            use_external_data_format=use_external_data_format,\n            one_external_file=one_external_file,\n        )\n\n        model_type = ORTConfigManager.get_model_ort_type(self.config.model_type)\n        optimization_options = optimization_config.create_fusion_options(model_type)\n\n        LOGGER.info(\"Optimizing model...\")\n\n        for model_path in self.onnx_model_path:\n            optimizer = optimize_model(\n                model_path.as_posix(),\n                model_type,\n                self.normalized_config.num_attention_heads,\n                self.normalized_config.hidden_size,\n                opt_level=optimization_config.optimization_level,\n                optimization_options=optimization_options,\n                use_gpu=optimization_config.optimize_for_gpu,\n                only_onnxruntime=not optimization_config.enable_transformers_specific_optimizations,\n            )\n\n            if optimization_config.fp16:\n                # keep_io_types to keep inputs/outputs as float32\n                optimizer.convert_float_to_float16(keep_io_types=True)\n\n            suffix = f\"_{file_suffix}\" if file_suffix else \"\"\n            output_path = save_dir.joinpath(f\"{model_path.stem}{suffix}\").with_suffix(model_path.suffix)\n            optimizer.save_model_to_file(output_path.as_posix(), use_external_data_format, one_external_file)\n\n        # Save the model configuration\n        self.config.save_pretrained(save_dir)\n        ort_config.save_pretrained(save_dir)\n\n        LOGGER.info(\n            f\"Optimized model saved at: {save_dir} (external data format: \"\n            f\"{use_external_data_format}; saved all tensor to one file: \"\n            f\"{one_external_file})\"\n        )\n\n        return Path(save_dir)\n\n    @staticmethod\n    def get_fused_operators(onnx_model_path: Union[str, os.PathLike]) -> Dict[str, int]:\n        \"\"\"\n        Computes the dictionary mapping the name of the fused operators to their number of apparition in the model.\n\n        Args:\n            onnx_model_path (`Union[str, os.PathLike]`):\n                Path of the ONNX model.\n\n        Returns:\n            The dictionary mapping the name of the fused operators to their number of apparition in the model.\n        \"\"\"\n        onnx_optimized_model = BertOnnxModel(load_model(onnx_model_path))\n        fused_operator = onnx_optimized_model.get_fused_operator_statistics()\n        LOGGER.info(\n            f\"The following operators were fused : { ', '.join([k for k,v in fused_operator.items() if v > 0])}\"\n        )\n        return {k: v for k, v in fused_operator.items() if v > 0}\n\n    @staticmethod\n    def get_nodes_number_difference(\n        onnx_model_path: Union[str, os.PathLike], onnx_optimized_model_path: Union[str, os.PathLike]\n    ) -> int:\n        \"\"\"\n        Compute the difference in the number of nodes between the original and the optimized model.\n\n        Args:\n            onnx_model_path (`Union[str, os.PathLike]`):\n                Path of the ONNX model.\n            onnx_optimized_model_path (`Union[str, os.PathLike]`):\n                Path of the optimized ONNX model.\n\n        Returns:\n            The difference in the number of nodes between the original and the optimized model.\n        \"\"\"\n        onnx_model = BertOnnxModel(load_model(onnx_model_path))\n        onnx_optimized_model = BertOnnxModel(load_model(onnx_optimized_model_path))\n\n        # Information in the number of nodes decrease resulting from optimization\n        nodes_number_onnx_model = len(onnx_model.nodes())\n        nodes_number_onnx_optimized_model = len(onnx_optimized_model.nodes())\n        difference_nodes_number = nodes_number_onnx_model - nodes_number_onnx_optimized_model\n        LOGGER.info(\n            f\"There are {nodes_number_onnx_model} nodes before optimization and {nodes_number_onnx_optimized_model}\"\n            f\"nodes after. The number of nodes removed is {difference_nodes_number}\"\n        )\n        return difference_nodes_number\n\n    @staticmethod\n    def get_operators_difference(\n        onnx_model_path: Union[str, os.PathLike], onnx_optimized_model_path: Union[str, os.PathLike]\n    ) -> Dict[str, int]:\n        \"\"\"\n        Compute the dictionary mapping the operators name to the difference in the number of corresponding nodes between\n        the original and the optimized model.\n\n        Args:\n            onnx_model_path (`Union[str, os.PathLike]`):\n                Path of the ONNX model.\n            onnx_optimized_model_path (`Union[str, os.PathLike]`):\n                Path of the optimized ONNX model.\n\n        Returns:\n            The dictionary mapping the operators name to the difference in the number of corresponding nodes between the\n            original and the optimized model.\n        \"\"\"\n        onnx_model = BertOnnxModel(load_model(onnx_model_path))\n        onnx_optimized_model = BertOnnxModel(load_model(onnx_optimized_model_path))\n\n        def nodes_difference_given_type(op_type):\n            onnx_model_nodes_with_op_type = len(onnx_model.get_nodes_by_op_type(op_type))\n            onnx_optimized_model_nodes_with_op_type = len(onnx_optimized_model.get_nodes_by_op_type(op_type))\n            return onnx_model_nodes_with_op_type - onnx_optimized_model_nodes_with_op_type\n\n        # Compute operators difference between the original and the optimized models\n        op_types = set()\n        for model in [onnx_model, onnx_optimized_model]:\n            for node in model.nodes():\n                op_types.add(node.op_type)\n\n        operators_difference = dict(map(lambda op_type: (op_type, nodes_difference_given_type(op_type)), op_types))\n        return {k: v for k, v in operators_difference.items() if v != 0}\n", "path": "optimum/onnxruntime/optimization.py"}], "after_files": [{"content": "#  Copyright 2021 The HuggingFace Team. All rights reserved.\n#\n#  Licensed under the Apache License, Version 2.0 (the \"License\");\n#  you may not use this file except in compliance with the License.\n#  You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n#  Unless required by applicable law or agreed to in writing, software\n#  distributed under the License is distributed on an \"AS IS\" BASIS,\n#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n#  See the License for the specific language governing permissions and\n#  limitations under the License.\n\"\"\"Main class for performing graph optimization with ONNX Runtime.\"\"\"\n\nimport logging\nimport os\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Dict, List, Optional, Union\n\nfrom transformers.models.auto.configuration_auto import AutoConfig\n\nfrom onnx import load_model\nfrom onnxruntime.transformers.onnx_model_bert import BertOnnxModel\nfrom onnxruntime.transformers.optimizer import optimize_model\n\nfrom ..utils import CONFIG_NAME, NormalizedConfigManager\nfrom ..utils.save_utils import maybe_save_preprocessors\nfrom .configuration import OptimizationConfig, ORTConfig\nfrom .modeling_ort import ORTModel\nfrom .modeling_seq2seq import ORTModelForSeq2SeqLM\nfrom .utils import ONNX_WEIGHTS_NAME, ORTConfigManager\n\n\nif TYPE_CHECKING:\n    from transformers import PretrainedConfig\n\n\nLOGGER = logging.getLogger(__name__)\n\n\nclass ORTOptimizer:\n    \"\"\"\n    Handles the ONNX Runtime optimization process for models shared on huggingface.co/models.\n    \"\"\"\n\n    def __init__(self, onnx_model_path: List[os.PathLike], config: \"PretrainedConfig\"):\n        \"\"\"\n        Args:\n            onnx_model_path (`List[os.PathLike]`):\n                The paths of the onnx models to optimize.\n            config ([`~transformers.PretrainedConfig`]):\n                An instance of the configuration associated to the model to optimize.\n        \"\"\"\n        super().__init__()\n        self.onnx_model_path = onnx_model_path\n        self.config = config\n        self.model_type = self.config.model_type\n        self.normalized_config = NormalizedConfigManager.get_normalized_config_class(self.model_type)(self.config)\n\n    @classmethod\n    def from_pretrained(\n        cls, model_or_path: Union[str, os.PathLike, ORTModel], file_names: Optional[List[str]] = None\n    ) -> \"ORTOptimizer\":\n        \"\"\"\n        Args:\n            model_or_path (`Union[str, os.PathLike, ORTModel]`):\n                The path to a local directory hosting the model to optimize or an instance of an `ORTModel` to quantize.\n                Can be either:\n                    - A path to a local *directory* containing the model to optimize.\n                    - An instance of [`~optimum.onnxruntime.ORTModel`].\n            file_names(`Optional[List[str]]`, *optional*):\n                The list of file names of the models to optimize.\n        \"\"\"\n        onnx_model_path = []\n        config = None\n        if isinstance(model_or_path, ORTModel):\n            if isinstance(model_or_path, ORTModelForSeq2SeqLM):\n                onnx_model_path += [\n                    model_or_path.encoder_model_path,\n                    model_or_path.decoder_model_path,\n                ]\n                # Add the decoder with past key/values if present\n                if model_or_path.use_cache:\n                    onnx_model_path.append(model_or_path.decoder_with_past_model_path)\n            else:\n                onnx_model_path.append(model_or_path.model_path)\n            config = model_or_path.config\n        elif os.path.isdir(model_or_path):\n            file_names = [ONNX_WEIGHTS_NAME] if file_names is None else file_names\n            model_or_path = Path(model_or_path)\n            if CONFIG_NAME not in os.listdir(model_or_path):\n                raise ValueError(f\"The local directory does not contain the configuration file {CONFIG_NAME}.\")\n            config = AutoConfig.from_pretrained(model_or_path)\n            for file_name in file_names:\n                onnx_model_path.append(model_or_path.joinpath(file_name))\n        else:\n            raise ValueError(f\"Unable to load the model from {model_or_path}.\")\n        return cls(onnx_model_path, config=config)\n\n    def optimize(\n        self,\n        optimization_config: OptimizationConfig,\n        save_dir: Union[str, os.PathLike],\n        file_suffix: Optional[str] = \"optimized\",\n        use_external_data_format: bool = False,\n        one_external_file: bool = True,\n    ):\n        \"\"\"\n        Optimizes a model given the optimization specifications defined in `optimization_config`.\n\n        Args:\n            optimization_config ([`~optimum.onnxruntime.OptimizationConfig`]):\n                The configuration containing the parameters related to optimization.\n            save_dir (`Union[str, os.PathLike]`):\n                The path used to save the optimized model.\n            file_suffix (`str`, *optional*, defaults to `\"optimized\"`):\n                The file suffix used to save the optimized model.\n            use_external_data_format (`bool`, *optional*, defaults to `False`):\n                Whether to use external data format to store model of size >= 2Gb.\n            one_external_file (`bool`, defaults to `True`):\n                When `use_external_data_format=True`, whether to save all tensors to one external file.\n                If false, save each tensor to a file named with the tensor name.\n\n        \"\"\"\n        save_dir = Path(save_dir)\n        save_dir.mkdir(parents=True, exist_ok=True)\n        ORTConfigManager.check_optimization_supported_model(self.model_type)\n\n        self.config.save_pretrained(save_dir)\n        maybe_save_preprocessors(self.onnx_model_path[0].parent, save_dir)\n\n        # Create and save the configuration summarizing all the parameters related to optimization\n        ort_config = ORTConfig(\n            optimization=optimization_config,\n            use_external_data_format=use_external_data_format,\n            one_external_file=one_external_file,\n        )\n\n        model_type = ORTConfigManager.get_model_ort_type(self.config.model_type)\n        optimization_options = optimization_config.create_fusion_options(model_type)\n\n        LOGGER.info(\"Optimizing model...\")\n\n        for model_path in self.onnx_model_path:\n            try:\n                optimizer = optimize_model(\n                    model_path.as_posix(),\n                    model_type,\n                    self.normalized_config.num_attention_heads,\n                    self.normalized_config.hidden_size,\n                    opt_level=optimization_config.optimization_level,\n                    optimization_options=optimization_options,\n                    use_gpu=optimization_config.optimize_for_gpu,\n                    only_onnxruntime=not optimization_config.enable_transformers_specific_optimizations,\n                )\n\n                if optimization_config.fp16:\n                    # keep_io_types to keep inputs/outputs as float32\n                    optimizer.convert_float_to_float16(\n                        use_symbolic_shape_infer=not optimization_config.disable_shape_inference, keep_io_types=True\n                    )\n            except Exception as e:\n                if \"Incomplete symbolic shape inference\" in str(e):\n                    err = RuntimeError(\n                        f\"{str(e)}. Try to set `disable_shape_inference=True` in your optimization configuration.\"\n                    )\n                    raise err from e\n                raise\n\n            suffix = f\"_{file_suffix}\" if file_suffix else \"\"\n            output_path = save_dir.joinpath(f\"{model_path.stem}{suffix}\").with_suffix(model_path.suffix)\n            optimizer.save_model_to_file(output_path.as_posix(), use_external_data_format, one_external_file)\n\n        # Save the model configuration\n        self.config.save_pretrained(save_dir)\n        ort_config.save_pretrained(save_dir)\n\n        LOGGER.info(\n            f\"Optimized model saved at: {save_dir} (external data format: \"\n            f\"{use_external_data_format}; saved all tensor to one file: \"\n            f\"{one_external_file})\"\n        )\n\n        return Path(save_dir)\n\n    @staticmethod\n    def get_fused_operators(onnx_model_path: Union[str, os.PathLike]) -> Dict[str, int]:\n        \"\"\"\n        Computes the dictionary mapping the name of the fused operators to their number of apparition in the model.\n\n        Args:\n            onnx_model_path (`Union[str, os.PathLike]`):\n                Path of the ONNX model.\n\n        Returns:\n            The dictionary mapping the name of the fused operators to their number of apparition in the model.\n        \"\"\"\n        onnx_optimized_model = BertOnnxModel(load_model(onnx_model_path))\n        fused_operator = onnx_optimized_model.get_fused_operator_statistics()\n        LOGGER.info(\n            f\"The following operators were fused : { ', '.join([k for k,v in fused_operator.items() if v > 0])}\"\n        )\n        return {k: v for k, v in fused_operator.items() if v > 0}\n\n    @staticmethod\n    def get_nodes_number_difference(\n        onnx_model_path: Union[str, os.PathLike], onnx_optimized_model_path: Union[str, os.PathLike]\n    ) -> int:\n        \"\"\"\n        Compute the difference in the number of nodes between the original and the optimized model.\n\n        Args:\n            onnx_model_path (`Union[str, os.PathLike]`):\n                Path of the ONNX model.\n            onnx_optimized_model_path (`Union[str, os.PathLike]`):\n                Path of the optimized ONNX model.\n\n        Returns:\n            The difference in the number of nodes between the original and the optimized model.\n        \"\"\"\n        onnx_model = BertOnnxModel(load_model(onnx_model_path))\n        onnx_optimized_model = BertOnnxModel(load_model(onnx_optimized_model_path))\n\n        # Information in the number of nodes decrease resulting from optimization\n        nodes_number_onnx_model = len(onnx_model.nodes())\n        nodes_number_onnx_optimized_model = len(onnx_optimized_model.nodes())\n        difference_nodes_number = nodes_number_onnx_model - nodes_number_onnx_optimized_model\n        LOGGER.info(\n            f\"There are {nodes_number_onnx_model} nodes before optimization and {nodes_number_onnx_optimized_model}\"\n            f\"nodes after. The number of nodes removed is {difference_nodes_number}\"\n        )\n        return difference_nodes_number\n\n    @staticmethod\n    def get_operators_difference(\n        onnx_model_path: Union[str, os.PathLike], onnx_optimized_model_path: Union[str, os.PathLike]\n    ) -> Dict[str, int]:\n        \"\"\"\n        Compute the dictionary mapping the operators name to the difference in the number of corresponding nodes between\n        the original and the optimized model.\n\n        Args:\n            onnx_model_path (`Union[str, os.PathLike]`):\n                Path of the ONNX model.\n            onnx_optimized_model_path (`Union[str, os.PathLike]`):\n                Path of the optimized ONNX model.\n\n        Returns:\n            The dictionary mapping the operators name to the difference in the number of corresponding nodes between the\n            original and the optimized model.\n        \"\"\"\n        onnx_model = BertOnnxModel(load_model(onnx_model_path))\n        onnx_optimized_model = BertOnnxModel(load_model(onnx_optimized_model_path))\n\n        def nodes_difference_given_type(op_type):\n            onnx_model_nodes_with_op_type = len(onnx_model.get_nodes_by_op_type(op_type))\n            onnx_optimized_model_nodes_with_op_type = len(onnx_optimized_model.get_nodes_by_op_type(op_type))\n            return onnx_model_nodes_with_op_type - onnx_optimized_model_nodes_with_op_type\n\n        # Compute operators difference between the original and the optimized models\n        op_types = set()\n        for model in [onnx_model, onnx_optimized_model]:\n            for node in model.nodes():\n                op_types.add(node.op_type)\n\n        operators_difference = dict(map(lambda op_type: (op_type, nodes_difference_given_type(op_type)), op_types))\n        return {k: v for k, v in operators_difference.items() if v != 0}\n", "path": "optimum/onnxruntime/optimization.py"}]}