We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add marian model type support for ORTOptimizer
### Feature request

The recent Optimum library enable to export marian model type to ONNX format. However, the optimization via ORTOptimizer seems not to support this model type.

Please find below, the code used:
```
from optimum.onnxruntime import ORTModelForSeq2SeqLM, ORTOptimizer
from optimum.onnxruntime.configuration import OptimizationConfig

model_id="Helsinki-NLP/opus-mt-fr-en"

# load vanilla transformers and convert to onnx
model = ORTModelForSeq2SeqLM.from_pretrained(model_id, from_transformers=True)

# Create ORTOptimizer
optimizer = ORTOptimizer.from_pretrained(model)

# Define the optimization strategy by creating the appropriate configuration
optimization_config = OptimizationConfig(optimization_level=1,
                                        optimize_for_gpu=True,
                                        fp16=True
                                        )

# Optimize the model
optimizer.optimize(save_dir=onnx_path, optimization_config=optimization_config)
```
And the error code on feedback:
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-8-0c22ffd4b88a> in <module>
     12 
     13 # Optimize the model
---> 14 optimizer.optimize(save_dir=onnx_path, optimization_config=optimization_config)

~/anaconda3/envs/sts-transformers-gpu-fresh/lib/python3.8/site-packages/optimum/onnxruntime/optimization.py in optimize(self, optimization_config, save_dir, file_suffix, use_external_data_format)
    113         save_dir.mkdir(parents=True, exist_ok=True)
    114         model_type = self.config.model_type
--> 115         ORTConfigManager.check_supported_model_or_raise(model_type)
    116 
    117         # Save the model configuration

~/anaconda3/envs/sts-transformers-gpu-fresh/lib/python3.8/site-packages/optimum/onnxruntime/utils.py in check_supported_model_or_raise(cls, model_type)
    110     def check_supported_model_or_raise(cls, model_type: str) -> bool:
    111         if model_type not in cls._conf:
--> 112             raise KeyError(
    113                 f"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. "
    114                 f"If you want to support {model_type} please propose a PR or open up an issue."

KeyError: "marian model type is not supported yet. Only ['bert', 'albert', 'big_bird', 'camembert', 'codegen', 'distilbert', 'deberta', 'deberta-v2', 'electra', 'roberta', 'bart', 'gpt2', 'gpt_neo', 'xlm-roberta'] are supported. If you want to support marian please propose a PR or open up an issue."
```

### Motivation

It would be greatly beneficial to be able to optimize the onnx models regarding the low inference speed of this kind of model.

### Your contribution

I could be a test user if the library is updated with this new feature.

--- 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/utils.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 from enum import Enum
15 
16 import torch
17 from transformers.onnx import OnnxConfig, OnnxSeq2SeqConfigWithPast
18 from transformers.utils import logging
19 
20 import onnx
21 import onnxruntime as ort
22 
23 from ..onnx import OnnxConfigWithLoss, OnnxSeq2SeqConfigWithPastAndLoss
24 
25 
26 logger = logging.get_logger(__name__)
27 
28 ONNX_WEIGHTS_NAME = "model.onnx"
29 OPTIMIZED_ONNX_WEIGHTS_NAME = "optimized_model.onnx"
30 QUANTIZED_ONNX_WEIGHTS_NAME = "q8_model.onnx"
31 
32 ONNX_ENCODER_NAME = "encoder_model.onnx"
33 ONNX_DECODER_NAME = "decoder_model.onnx"
34 ONNX_DECODER_WITH_PAST_NAME = "decoder_with_past_model.onnx"
35 
36 
37 def _is_gpu_available():
38     """
39     checks if a gpu is available.
40     """
41     available_providers = ort.get_available_providers()
42     if "CUDAExecutionProvider" in available_providers and torch.cuda.is_available():
43         return True
44     else:
45         return False
46 
47 
48 class ORTConfigManager:
49     """
50     A class that contains all the information needed by ONNX Runtime optimization for a given model type.
51 
52     Attributes:
53         _conf (`Dict[str, tuple]`):
54             A dictionary mapping each supported model type to a tuple containing the number of attention heads
55             and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.
56     """
57 
58     _conf = {
59         "bert": ("num_attention_heads", "hidden_size", "bert"),
60         "albert": ("num_attention_heads", "hidden_size", "bert"),
61         "big_bird": ("num_attention_heads", "hidden_size", "bert"),
62         "camembert": ("num_attention_heads", "hidden_size", "bert"),
63         "codegen": ("n_head", "n_embd", "gpt2"),
64         "distilbert": ("n_heads", "dim", "bert"),
65         "deberta": ("num_attention_heads", "hidden_size", "bert"),
66         "deberta-v2": ("num_attention_heads", "hidden_size", "bert"),
67         "electra": ("num_attention_heads", "hidden_size", "bert"),
68         "roberta": ("num_attention_heads", "hidden_size", "bert"),
69         "bart": ("encoder_attention_heads", "d_model", "bart"),
70         "gpt2": ("n_head", "n_embd", "gpt2"),
71         "gpt_neo": ("num_heads", "hidden_size", "gpt2"),
72         "xlm-roberta": ("num_attention_heads", "hidden_size", "bert"),
73     }
74 
75     @classmethod
76     def get_num_heads_name(cls, model_type: str) -> str:
77         num_heads = "num_attention_heads"
78         try:
79             num_heads = cls._conf[model_type][0]
80         except KeyError:
81             logger.warning(
82                 f"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to "
83                 f"access the number of heads defined in the config is set to `{num_heads}`."
84             )
85         return num_heads
86 
87     @classmethod
88     def get_hidden_size_name(cls, model_type: str) -> str:
89         hidden_size = "hidden_size"
90         try:
91             hidden_size = cls._conf[model_type][1]
92         except KeyError:
93             logger.warning(
94                 f"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to "
95                 f"access the hidden size defined in the config is set to `{hidden_size}`."
96             )
97         return hidden_size
98 
99     @classmethod
100     def get_model_ort_type(cls, model_type: str) -> str:
101         try:
102             model_type = cls._conf[model_type][2]
103         except KeyError:
104             logger.warning(f"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported.")
105         return model_type
106 
107     @classmethod
108     def check_supported_model_or_raise(cls, model_type: str) -> bool:
109         if model_type not in cls._conf:
110             raise KeyError(
111                 f"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. "
112                 f"If you want to support {model_type} please propose a PR or open up an issue."
113             )
114 
115 
116 def generate_identified_filename(filename, identifier):
117     return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix)
118 
119 
120 def fix_atenops_to_gather(model_path):
121     # Fix broken ATenOp nodes back to Gather nodes.
122     model = onnx.load(model_path)
123     onnx.checker.check_model(model)
124 
125     nodes = model.graph.node
126 
127     for node in nodes:
128         if node.op_type in ["ATenOp", "ATen"]:
129             logger.info(f"----Start fixing node: {node.name}----")
130             op_num = node.name.split("_")[-1]
131             new_node = onnx.helper.make_node(
132                 "Gather",
133                 name="Gather_" + op_num,
134                 inputs=[node.input[0], node.input[1]],
135                 outputs=node.output,
136             )
137 
138             model.graph.node.remove(node)
139             model.graph.node.insert(int(op_num), new_node)
140 
141     onnx.checker.check_model(model)
142     onnx.save(model, model_path)
143 
144 
145 def wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:
146     if isinstance(onnx_config, OnnxSeq2SeqConfigWithPast):
147         return OnnxSeq2SeqConfigWithPastAndLoss(onnx_config)
148     else:
149         return OnnxConfigWithLoss(onnx_config)
150 
151 
152 def get_device_for_provider(provider: str) -> torch.device:
153     """
154     Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.
155     """
156     return torch.device("cuda") if provider == "CUDAExecutionProvider" else torch.device("cpu")
157 
158 
159 def get_provider_for_device(device: torch.device) -> str:
160     """
161     Gets the ONNX Runtime provider associated with the PyTorch device (CPU/CUDA).
162     """
163     return "CUDAExecutionProvider" if device.type.lower() == "cuda" else "CPUExecutionProvider"
164 
165 
166 class ORTQuantizableOperator(Enum):
167     # Common ops
168     Gather = "Gather"
169     Transpose = "Transpose"
170     EmbedLayerNormalizationQuant = "EmbedLayerNormalization"
171 
172     # QLinearOps
173     Conv = "Conv"
174     MatMul = "MatMul"
175     Add = "Add"
176     Mul = "Mul"
177     Relu = "Relu"
178     Clip = "Clip"
179     LeakyRelu = "LeakyRelu"
180     Sigmoid = "Sigmoid"
181     MaxPool = "MaxPool"
182     GlobalAveragePool = "GlobalAveragePool"
183     Split = "Split"
184     Pad = "Pad"
185     Reshape = "Reshape"
186     Squeeze = "Squeeze"
187     Unsqueeze = "Unsqueeze"
188     Resize = "Resize"
189     AveragePool = "AveragePool"
190     Concat = "Concat"
191
```
--- 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/utils.py b/optimum/onnxruntime/utils.py
--- a/optimum/onnxruntime/utils.py
+++ b/optimum/onnxruntime/utils.py
@@ -55,20 +55,22 @@
             and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.
     """
 
+    # Contribution note: Please add new models in alphabetical order
     _conf = {
-        "bert": ("num_attention_heads", "hidden_size", "bert"),
         "albert": ("num_attention_heads", "hidden_size", "bert"),
+        "bart": ("encoder_attention_heads", "d_model", "bart"),
+        "bert": ("num_attention_heads", "hidden_size", "bert"),
         "big_bird": ("num_attention_heads", "hidden_size", "bert"),
         "camembert": ("num_attention_heads", "hidden_size", "bert"),
         "codegen": ("n_head", "n_embd", "gpt2"),
-        "distilbert": ("n_heads", "dim", "bert"),
         "deberta": ("num_attention_heads", "hidden_size", "bert"),
         "deberta-v2": ("num_attention_heads", "hidden_size", "bert"),
+        "distilbert": ("n_heads", "dim", "bert"),
         "electra": ("num_attention_heads", "hidden_size", "bert"),
-        "roberta": ("num_attention_heads", "hidden_size", "bert"),
-        "bart": ("encoder_attention_heads", "d_model", "bart"),
         "gpt2": ("n_head", "n_embd", "gpt2"),
         "gpt_neo": ("num_heads", "hidden_size", "gpt2"),
+        "marian": ("encoder_attention_heads", "d_model", "bart"),
+        "roberta": ("num_attention_heads", "hidden_size", "bert"),
         "xlm-roberta": ("num_attention_heads", "hidden_size", "bert"),
     }

{"golden_diff": "diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py\n--- a/optimum/onnxruntime/utils.py\n+++ b/optimum/onnxruntime/utils.py\n@@ -55,20 +55,22 @@\n             and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.\n     \"\"\"\n \n+    # Contribution note: Please add new models in alphabetical order\n     _conf = {\n-        \"bert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"albert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n+        \"bart\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n+        \"bert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"big_bird\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"camembert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"codegen\": (\"n_head\", \"n_embd\", \"gpt2\"),\n-        \"distilbert\": (\"n_heads\", \"dim\", \"bert\"),\n         \"deberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"deberta-v2\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n+        \"distilbert\": (\"n_heads\", \"dim\", \"bert\"),\n         \"electra\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n-        \"roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n-        \"bart\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n         \"gpt2\": (\"n_head\", \"n_embd\", \"gpt2\"),\n         \"gpt_neo\": (\"num_heads\", \"hidden_size\", \"gpt2\"),\n+        \"marian\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n+        \"roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n         \"xlm-roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n     }\n", "issue": "Add marian model type support for ORTOptimizer\n### Feature request\n\nThe recent Optimum library enable to export marian model type to ONNX format. However, the optimization via ORTOptimizer seems not to support this model type.\r\n\r\nPlease find below, the code used:\r\n```\r\nfrom optimum.onnxruntime import ORTModelForSeq2SeqLM, ORTOptimizer\r\nfrom optimum.onnxruntime.configuration import OptimizationConfig\r\n\r\nmodel_id=\"Helsinki-NLP/opus-mt-fr-en\"\r\n\r\n# load vanilla transformers and convert to onnx\r\nmodel = ORTModelForSeq2SeqLM.from_pretrained(model_id, from_transformers=True)\r\n\r\n# Create ORTOptimizer\r\noptimizer = ORTOptimizer.from_pretrained(model)\r\n\r\n# Define the optimization strategy by creating the appropriate configuration\r\noptimization_config = OptimizationConfig(optimization_level=1,\r\n                                        optimize_for_gpu=True,\r\n                                        fp16=True\r\n                                        )\r\n\r\n# Optimize the model\r\noptimizer.optimize(save_dir=onnx_path, optimization_config=optimization_config)\r\n```\r\nAnd the error code on feedback:\r\n```\r\n---------------------------------------------------------------------------\r\nKeyError                                  Traceback (most recent call last)\r\n<ipython-input-8-0c22ffd4b88a> in <module>\r\n     12 \r\n     13 # Optimize the model\r\n---> 14 optimizer.optimize(save_dir=onnx_path, optimization_config=optimization_config)\r\n\r\n~/anaconda3/envs/sts-transformers-gpu-fresh/lib/python3.8/site-packages/optimum/onnxruntime/optimization.py in optimize(self, optimization_config, save_dir, file_suffix, use_external_data_format)\r\n    113         save_dir.mkdir(parents=True, exist_ok=True)\r\n    114         model_type = self.config.model_type\r\n--> 115         ORTConfigManager.check_supported_model_or_raise(model_type)\r\n    116 \r\n    117         # Save the model configuration\r\n\r\n~/anaconda3/envs/sts-transformers-gpu-fresh/lib/python3.8/site-packages/optimum/onnxruntime/utils.py in check_supported_model_or_raise(cls, model_type)\r\n    110     def check_supported_model_or_raise(cls, model_type: str) -> bool:\r\n    111         if model_type not in cls._conf:\r\n--> 112             raise KeyError(\r\n    113                 f\"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. \"\r\n    114                 f\"If you want to support {model_type} please propose a PR or open up an issue.\"\r\n\r\nKeyError: \"marian model type is not supported yet. Only ['bert', 'albert', 'big_bird', 'camembert', 'codegen', 'distilbert', 'deberta', 'deberta-v2', 'electra', 'roberta', 'bart', 'gpt2', 'gpt_neo', 'xlm-roberta'] are supported. If you want to support marian please propose a PR or open up an issue.\"\r\n```\n\n### Motivation\n\nIt would be greatly beneficial to be able to optimize the onnx models regarding the low inference speed of this kind of model.\n\n### Your contribution\n\nI could be a test user if the library is updated with this new feature.\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.\nfrom enum import Enum\n\nimport torch\nfrom transformers.onnx import OnnxConfig, OnnxSeq2SeqConfigWithPast\nfrom transformers.utils import logging\n\nimport onnx\nimport onnxruntime as ort\n\nfrom ..onnx import OnnxConfigWithLoss, OnnxSeq2SeqConfigWithPastAndLoss\n\n\nlogger = logging.get_logger(__name__)\n\nONNX_WEIGHTS_NAME = \"model.onnx\"\nOPTIMIZED_ONNX_WEIGHTS_NAME = \"optimized_model.onnx\"\nQUANTIZED_ONNX_WEIGHTS_NAME = \"q8_model.onnx\"\n\nONNX_ENCODER_NAME = \"encoder_model.onnx\"\nONNX_DECODER_NAME = \"decoder_model.onnx\"\nONNX_DECODER_WITH_PAST_NAME = \"decoder_with_past_model.onnx\"\n\n\ndef _is_gpu_available():\n    \"\"\"\n    checks if a gpu is available.\n    \"\"\"\n    available_providers = ort.get_available_providers()\n    if \"CUDAExecutionProvider\" in available_providers and torch.cuda.is_available():\n        return True\n    else:\n        return False\n\n\nclass ORTConfigManager:\n    \"\"\"\n    A class that contains all the information needed by ONNX Runtime optimization for a given model type.\n\n    Attributes:\n        _conf (`Dict[str, tuple]`):\n            A dictionary mapping each supported model type to a tuple containing the number of attention heads\n            and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.\n    \"\"\"\n\n    _conf = {\n        \"bert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"albert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"big_bird\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"camembert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"codegen\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"distilbert\": (\"n_heads\", \"dim\", \"bert\"),\n        \"deberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"deberta-v2\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"electra\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"bart\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n        \"gpt2\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"gpt_neo\": (\"num_heads\", \"hidden_size\", \"gpt2\"),\n        \"xlm-roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n    }\n\n    @classmethod\n    def get_num_heads_name(cls, model_type: str) -> str:\n        num_heads = \"num_attention_heads\"\n        try:\n            num_heads = cls._conf[model_type][0]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the number of heads defined in the config is set to `{num_heads}`.\"\n            )\n        return num_heads\n\n    @classmethod\n    def get_hidden_size_name(cls, model_type: str) -> str:\n        hidden_size = \"hidden_size\"\n        try:\n            hidden_size = cls._conf[model_type][1]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the hidden size defined in the config is set to `{hidden_size}`.\"\n            )\n        return hidden_size\n\n    @classmethod\n    def get_model_ort_type(cls, model_type: str) -> str:\n        try:\n            model_type = cls._conf[model_type][2]\n        except KeyError:\n            logger.warning(f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported.\")\n        return model_type\n\n    @classmethod\n    def check_supported_model_or_raise(cls, model_type: str) -> bool:\n        if model_type not in cls._conf:\n            raise KeyError(\n                f\"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. \"\n                f\"If you want to support {model_type} please propose a PR or open up an issue.\"\n            )\n\n\ndef generate_identified_filename(filename, identifier):\n    return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix)\n\n\ndef fix_atenops_to_gather(model_path):\n    # Fix broken ATenOp nodes back to Gather nodes.\n    model = onnx.load(model_path)\n    onnx.checker.check_model(model)\n\n    nodes = model.graph.node\n\n    for node in nodes:\n        if node.op_type in [\"ATenOp\", \"ATen\"]:\n            logger.info(f\"----Start fixing node: {node.name}----\")\n            op_num = node.name.split(\"_\")[-1]\n            new_node = onnx.helper.make_node(\n                \"Gather\",\n                name=\"Gather_\" + op_num,\n                inputs=[node.input[0], node.input[1]],\n                outputs=node.output,\n            )\n\n            model.graph.node.remove(node)\n            model.graph.node.insert(int(op_num), new_node)\n\n    onnx.checker.check_model(model)\n    onnx.save(model, model_path)\n\n\ndef wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:\n    if isinstance(onnx_config, OnnxSeq2SeqConfigWithPast):\n        return OnnxSeq2SeqConfigWithPastAndLoss(onnx_config)\n    else:\n        return OnnxConfigWithLoss(onnx_config)\n\n\ndef get_device_for_provider(provider: str) -> torch.device:\n    \"\"\"\n    Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.\n    \"\"\"\n    return torch.device(\"cuda\") if provider == \"CUDAExecutionProvider\" else torch.device(\"cpu\")\n\n\ndef get_provider_for_device(device: torch.device) -> str:\n    \"\"\"\n    Gets the ONNX Runtime provider associated with the PyTorch device (CPU/CUDA).\n    \"\"\"\n    return \"CUDAExecutionProvider\" if device.type.lower() == \"cuda\" else \"CPUExecutionProvider\"\n\n\nclass ORTQuantizableOperator(Enum):\n    # Common ops\n    Gather = \"Gather\"\n    Transpose = \"Transpose\"\n    EmbedLayerNormalizationQuant = \"EmbedLayerNormalization\"\n\n    # QLinearOps\n    Conv = \"Conv\"\n    MatMul = \"MatMul\"\n    Add = \"Add\"\n    Mul = \"Mul\"\n    Relu = \"Relu\"\n    Clip = \"Clip\"\n    LeakyRelu = \"LeakyRelu\"\n    Sigmoid = \"Sigmoid\"\n    MaxPool = \"MaxPool\"\n    GlobalAveragePool = \"GlobalAveragePool\"\n    Split = \"Split\"\n    Pad = \"Pad\"\n    Reshape = \"Reshape\"\n    Squeeze = \"Squeeze\"\n    Unsqueeze = \"Unsqueeze\"\n    Resize = \"Resize\"\n    AveragePool = \"AveragePool\"\n    Concat = \"Concat\"\n", "path": "optimum/onnxruntime/utils.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.\nfrom enum import Enum\n\nimport torch\nfrom transformers.onnx import OnnxConfig, OnnxSeq2SeqConfigWithPast\nfrom transformers.utils import logging\n\nimport onnx\nimport onnxruntime as ort\n\nfrom ..onnx import OnnxConfigWithLoss, OnnxSeq2SeqConfigWithPastAndLoss\n\n\nlogger = logging.get_logger(__name__)\n\nONNX_WEIGHTS_NAME = \"model.onnx\"\nOPTIMIZED_ONNX_WEIGHTS_NAME = \"optimized_model.onnx\"\nQUANTIZED_ONNX_WEIGHTS_NAME = \"q8_model.onnx\"\n\nONNX_ENCODER_NAME = \"encoder_model.onnx\"\nONNX_DECODER_NAME = \"decoder_model.onnx\"\nONNX_DECODER_WITH_PAST_NAME = \"decoder_with_past_model.onnx\"\n\n\ndef _is_gpu_available():\n    \"\"\"\n    checks if a gpu is available.\n    \"\"\"\n    available_providers = ort.get_available_providers()\n    if \"CUDAExecutionProvider\" in available_providers and torch.cuda.is_available():\n        return True\n    else:\n        return False\n\n\nclass ORTConfigManager:\n    \"\"\"\n    A class that contains all the information needed by ONNX Runtime optimization for a given model type.\n\n    Attributes:\n        _conf (`Dict[str, tuple]`):\n            A dictionary mapping each supported model type to a tuple containing the number of attention heads\n            and the hidden size model config attribute names as well as the corresponding ONNX Runtime model type.\n    \"\"\"\n\n    # Contribution note: Please add new models in alphabetical order\n    _conf = {\n        \"albert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"bart\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n        \"bert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"big_bird\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"camembert\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"codegen\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"deberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"deberta-v2\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"distilbert\": (\"n_heads\", \"dim\", \"bert\"),\n        \"electra\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"gpt2\": (\"n_head\", \"n_embd\", \"gpt2\"),\n        \"gpt_neo\": (\"num_heads\", \"hidden_size\", \"gpt2\"),\n        \"marian\": (\"encoder_attention_heads\", \"d_model\", \"bart\"),\n        \"roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n        \"xlm-roberta\": (\"num_attention_heads\", \"hidden_size\", \"bert\"),\n    }\n\n    @classmethod\n    def get_num_heads_name(cls, model_type: str) -> str:\n        num_heads = \"num_attention_heads\"\n        try:\n            num_heads = cls._conf[model_type][0]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the number of heads defined in the config is set to `{num_heads}`.\"\n            )\n        return num_heads\n\n    @classmethod\n    def get_hidden_size_name(cls, model_type: str) -> str:\n        hidden_size = \"hidden_size\"\n        try:\n            hidden_size = cls._conf[model_type][1]\n        except KeyError:\n            logger.warning(\n                f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported. The default value to \"\n                f\"access the hidden size defined in the config is set to `{hidden_size}`.\"\n            )\n        return hidden_size\n\n    @classmethod\n    def get_model_ort_type(cls, model_type: str) -> str:\n        try:\n            model_type = cls._conf[model_type][2]\n        except KeyError:\n            logger.warning(f\"{model_type} is not supported yet. Only {list(cls._conf.keys())} are supported.\")\n        return model_type\n\n    @classmethod\n    def check_supported_model_or_raise(cls, model_type: str) -> bool:\n        if model_type not in cls._conf:\n            raise KeyError(\n                f\"{model_type} model type is not supported yet. Only {list(cls._conf.keys())} are supported. \"\n                f\"If you want to support {model_type} please propose a PR or open up an issue.\"\n            )\n\n\ndef generate_identified_filename(filename, identifier):\n    return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix)\n\n\ndef fix_atenops_to_gather(model_path):\n    # Fix broken ATenOp nodes back to Gather nodes.\n    model = onnx.load(model_path)\n    onnx.checker.check_model(model)\n\n    nodes = model.graph.node\n\n    for node in nodes:\n        if node.op_type in [\"ATenOp\", \"ATen\"]:\n            logger.info(f\"----Start fixing node: {node.name}----\")\n            op_num = node.name.split(\"_\")[-1]\n            new_node = onnx.helper.make_node(\n                \"Gather\",\n                name=\"Gather_\" + op_num,\n                inputs=[node.input[0], node.input[1]],\n                outputs=node.output,\n            )\n\n            model.graph.node.remove(node)\n            model.graph.node.insert(int(op_num), new_node)\n\n    onnx.checker.check_model(model)\n    onnx.save(model, model_path)\n\n\ndef wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:\n    if isinstance(onnx_config, OnnxSeq2SeqConfigWithPast):\n        return OnnxSeq2SeqConfigWithPastAndLoss(onnx_config)\n    else:\n        return OnnxConfigWithLoss(onnx_config)\n\n\ndef get_device_for_provider(provider: str) -> torch.device:\n    \"\"\"\n    Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.\n    \"\"\"\n    return torch.device(\"cuda\") if provider == \"CUDAExecutionProvider\" else torch.device(\"cpu\")\n\n\ndef get_provider_for_device(device: torch.device) -> str:\n    \"\"\"\n    Gets the ONNX Runtime provider associated with the PyTorch device (CPU/CUDA).\n    \"\"\"\n    return \"CUDAExecutionProvider\" if device.type.lower() == \"cuda\" else \"CPUExecutionProvider\"\n\n\nclass ORTQuantizableOperator(Enum):\n    # Common ops\n    Gather = \"Gather\"\n    Transpose = \"Transpose\"\n    EmbedLayerNormalizationQuant = \"EmbedLayerNormalization\"\n\n    # QLinearOps\n    Conv = \"Conv\"\n    MatMul = \"MatMul\"\n    Add = \"Add\"\n    Mul = \"Mul\"\n    Relu = \"Relu\"\n    Clip = \"Clip\"\n    LeakyRelu = \"LeakyRelu\"\n    Sigmoid = \"Sigmoid\"\n    MaxPool = \"MaxPool\"\n    GlobalAveragePool = \"GlobalAveragePool\"\n    Split = \"Split\"\n    Pad = \"Pad\"\n    Reshape = \"Reshape\"\n    Squeeze = \"Squeeze\"\n    Unsqueeze = \"Unsqueeze\"\n    Resize = \"Resize\"\n    AveragePool = \"AveragePool\"\n    Concat = \"Concat\"\n", "path": "optimum/onnxruntime/utils.py"}]}