We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
`*with_loss` wrappers failed for transformers 4.19.0
# Problem
As new the new argument `num_choices` added in `transformers.onnx.config.generate_dummy_inputs`, the wrappers will fail to match correct arguments

### Error Message
------------
```
Traceback (most recent call last):
  File "test_onnxruntime_train.py", line 116, in test_ort_trainer
    ort_eval_metrics = trainer.evaluate(inference_with_ort=inference_with_ort)
  File "/workspace/optimum/onnxruntime/trainer.py", line 631, in evaluate
    output = eval_loop(
  File "/workspace/optimum/onnxruntime/trainer.py", line 767, in evaluation_loop_ort
    self._export(onnx_model_path, with_loss=with_loss)
  File "/workspace/optimum/onnxruntime/trainer.py", line 1230, in _export
    _ = export(preprocessor=self.tokenizer, model=model, config=onnx_config, opset=opset, output=model_path)
  File "/usr/local/lib/python3.8/dist-packages/transformers/onnx/convert.py", line 313, in export
    return export_pytorch(preprocessor, model, config, opset, output, tokenizer=tokenizer)
  File "/usr/local/lib/python3.8/dist-packages/transformers/onnx/convert.py", line 138, in export_pytorch
    model_inputs = config.generate_dummy_inputs(preprocessor, framework=TensorType.PYTORCH)
  File "/workspace/optimum/onnx/configuration.py", line 170, in generate_dummy_inputs
    dummy_inputs = super().generate_dummy_inputs(
  File "/usr/local/lib/python3.8/dist-packages/transformers/onnx/config.py", line 308, in generate_dummy_inputs
    token_to_add = preprocessor.num_special_tokens_to_add(is_pair)
  File "/usr/local/lib/python3.8/dist-packages/transformers/tokenization_utils_fast.py", line 289, in num_special_tokens_to_add
    return self._tokenizer.num_special_tokens_to_add(pair)
TypeError: Can't convert <TensorType.PYTORCH: 'pt'> to PyBool
```

# Solution
Add keys for arguments

--- 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/onnx/configuration.py`
Content:
```
1 #  Copyright 2022 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 
15 import copy
16 from abc import ABC
17 from collections import OrderedDict
18 from typing import TYPE_CHECKING, Any, Mapping, Optional, Union
19 
20 from transformers.file_utils import TensorType, is_tf_available, is_torch_available
21 from transformers.onnx.utils import compute_effective_axis_dimension
22 from transformers.utils import logging
23 
24 
25 if TYPE_CHECKING:
26     from transformers.configuration_utils import PretrainedConfig
27     from transformers.feature_extraction_utils import FeatureExtractionMixin
28     from transformers.tokenization_utils_base import PreTrainedTokenizerBase
29 
30 from transformers.onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast
31 
32 
33 logger = logging.get_logger(__name__)
34 
35 
36 class OnnxConfigWithLoss(OnnxConfig, ABC):
37     """
38     Wrapper for the children classes of `transformers.onnx.OnnxConfig` to export the model through the ONNX format with loss in outputs.
39     """
40 
41     _tasks_to_extra_inputs = {
42         "default": OrderedDict({"labels": {0: "batch"}}),
43         "masked-lm": OrderedDict({"labels": {0: "batch", 1: "sequence"}}),
44         "causal-lm": OrderedDict({"labels": {0: "batch", 1: "sequence"}}),
45         "seq2seq-lm": OrderedDict({"labels": {0: "batch", 1: "sequence"}}),
46         "sequence-classification": OrderedDict({"labels": {0: "batch"}}),
47         "token-classification": OrderedDict({"labels": {0: "batch", 1: "sequence"}}),
48         "multiple-choice": OrderedDict({"labels": {0: "batch"}}),
49         "question-answering": OrderedDict(
50             {
51                 "start_positions": {0: "batch"},
52                 "end_positions": {0: "batch"},
53             }
54         ),
55         "image-classification": OrderedDict({"labels": {0: "batch"}}),
56     }
57     _tasks_to_extra_outputs = {
58         "default": OrderedDict({"loss": {}}),
59     }
60 
61     def __init__(self, config: OnnxConfig):
62         self.__dict__ = copy.deepcopy(config.__dict__)
63         self._onnx_config = config
64         if self.task not in self._tasks_to_extra_inputs:
65             raise ValueError(
66                 f"{self.task} is not a supported task, supported tasks: {self._tasks_to_extra_inputs.keys()}"
67             )
68 
69     @classmethod
70     def from_model_config(cls, config: OnnxConfig) -> "OnnxConfigWithLoss":
71         """
72         Instantiate a OnnxConfigWithLoss for a specific model
73         Args:
74             config: The model's configuration to use when exporting to ONNX
75         Returns:
76             OnnxConfigWithLoss for this model
77         """
78         return cls(config)
79 
80     @property
81     def inputs(self) -> Mapping[str, Mapping[int, str]]:
82         """
83         Mapping containing the axis definition of the input tensors(including labels) to provide to the model
84         Returns:
85             For each input: its name associated to the axes symbolic name and the axis position within the tensor
86         """
87         inputs = self._onnx_config.inputs
88         inputs.update(self._tasks_to_extra_inputs[self.task])
89         return inputs
90 
91     @property
92     def outputs(self) -> Mapping[str, Mapping[int, str]]:
93         """
94         Mapping containing the axis definition of the output tensors(including loss) to provide to the model
95         Returns:
96             For each output: its name associated to the axes symbolic name and the axis position within the tensor
97         """
98         common_outputs = self._onnx_config.outputs
99         extra_outputs = self._tasks_to_extra_outputs["default"]
100         common_outputs.update(extra_outputs)
101         for key in reversed(extra_outputs.keys()):
102             common_outputs.move_to_end(key, last=False)
103         return copy.deepcopy(common_outputs)
104 
105     def _generate_extra_dummy_inputs_pt(
106         self,
107         dummy_inputs,
108         batch_size,
109         seq_length,
110     ) -> Mapping[str, Any]:
111         import torch
112 
113         for label, input in self._tasks_to_extra_inputs[self.task].items():
114             if "sequence" in input.values():
115                 dummy_inputs[label] = torch.zeros(batch_size, seq_length, dtype=torch.long)
116             else:
117                 dummy_inputs[label] = torch.zeros(batch_size, dtype=torch.long)
118         return dummy_inputs
119 
120     def _generate_extra_dummy_inputs_tf(
121         self,
122         dummy_inputs,
123         batch_size,
124         seq_length,
125     ) -> Mapping[str, Any]:
126         import tensorflow as tf
127 
128         for label, input in self._tasks_to_extra_inputs[self.task].items():
129             if "sequence" in input.values():
130                 dummy_inputs[label] = tf.zeros(batch_size, seq_length, dtype=tf.int64)
131             else:
132                 dummy_inputs[label] = tf.zeros(batch_size, dtype=tf.int64)
133         return dummy_inputs
134 
135     def generate_dummy_inputs(
136         self,
137         preprocessor: Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"],
138         batch_size: int = -1,
139         seq_length: int = -1,
140         is_pair: bool = False,
141         framework: Optional[TensorType] = None,
142         num_channels: int = 3,
143         image_width: int = 40,
144         image_height: int = 40,
145         tokenizer: "PreTrainedTokenizerBase" = None,
146     ) -> Mapping[str, Any]:
147         """
148         Generate inputs to provide to the ONNX exporter for the specific framework
149         Args:
150             preprocessor: ([`PreTrainedTokenizerBase`] or [`FeatureExtractionMixin`]):
151                 The preprocessor associated with this model configuration.
152             batch_size (`int`, *optional*, defaults to -1):
153                 The batch size to export the model for (-1 means dynamic axis).
154             seq_length (`int`, *optional*, defaults to -1):
155                 The sequence length to export the model for (-1 means dynamic axis).
156             is_pair (`bool`, *optional*, defaults to `False`):
157                 Indicate if the input is a pair (sentence 1, sentence 2)
158             framework (`TensorType`, *optional*, defaults to `None`):
159                 The framework (PyTorch or TensorFlow) that the tokenizer will generate tensors for.
160             num_channels (`int`, *optional*, defaults to 3):
161                 The number of channels of the generated images.
162             image_width (`int`, *optional*, defaults to 40):
163                 The width of the generated images.
164             image_height (`int`, *optional*, defaults to 40):
165                 The height of the generated images.
166         Returns:
167             Mapping[str, Tensor] holding the kwargs to provide to the model's forward function
168         """
169         # Generate dummy labels
170         dummy_inputs = super().generate_dummy_inputs(
171             preprocessor,
172             batch_size,
173             seq_length,
174             is_pair,
175             framework,
176             num_channels,
177             image_width,
178             image_height,
179             tokenizer,
180         )
181         label_batch_size = compute_effective_axis_dimension(
182             batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0
183         )
184         label_seq_length = compute_effective_axis_dimension(
185             seq_length, fixed_dimension=self.default_fixed_sequence, num_token_to_add=0
186         )
187 
188         if framework == TensorType.PYTORCH:
189             if is_torch_available():
190                 return self._generate_extra_dummy_inputs_pt(dummy_inputs, label_batch_size, label_seq_length)
191             else:
192                 raise RuntimeError(f"Could not generate dummy inputs because no PyTorch installation was found.")
193         elif framework == TensorType.TENSORFLOW:
194             if is_tf_available():
195                 return self._generate_extra_dummy_inputs_tf(dummy_inputs, label_batch_size, label_seq_length)
196             else:
197                 raise RuntimeError(f"Could not generate dummy inputs because no TensorFlow installation was found.")
198         else:
199             raise ValueError(
200                 f"Only two frameworks are supported for ONNX export: PyTorch or TensorFlow, but {framework} was provided."
201             )
202 
203 
204 class OnnxConfigWithPastAndLoss(OnnxConfigWithLoss, ABC):
205     def __init__(
206         self,
207         config: OnnxConfigWithPast,
208         use_past: bool = False,
209     ):
210         super().__init__(config)
211         self.use_past = use_past
212 
213     @classmethod
214     def with_past(cls, config: OnnxConfigWithPast) -> "OnnxConfigWithPast":
215         """
216         Instantiate a OnnxConfigWithPastAndLoss with `use_past` attribute set to True
217         Args:
218             config: The underlying model's config to use when exporting to ONNX
219         Returns:
220             OnnxConfigWithPastAndLoss with `.use_past = True`
221         """
222         return cls(config, use_past=True)
223 
224     @property
225     def outputs(self) -> Mapping[str, Mapping[int, str]]:
226         common_outputs = super().outputs
227         if self.use_past:
228             self._onnx_config.fill_with_past_key_values_(common_outputs, direction="outputs")
229 
230         return common_outputs
231 
232     def generate_dummy_inputs(
233         self,
234         tokenizer: "PreTrainedTokenizerBase",
235         batch_size: int = -1,
236         seq_length: int = -1,
237         is_pair: bool = False,
238         framework: Optional[TensorType] = None,
239     ) -> Mapping[str, Any]:
240 
241         dummy_inputs = self._onnx_config.generate_dummy_inputs(
242             tokenizer,
243             batch_size,
244             seq_length,
245             is_pair,
246             framework,
247         )
248         label_batch_size = compute_effective_axis_dimension(
249             batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0
250         )
251         label_seq_length = compute_effective_axis_dimension(
252             seq_length, fixed_dimension=self.default_fixed_sequence, num_token_to_add=0
253         )
254 
255         if framework == TensorType.PYTORCH:
256             if is_torch_available():
257                 return self._generate_extra_dummy_inputs_pt(dummy_inputs, label_batch_size, label_seq_length)
258             else:
259                 raise RuntimeError(f"Could not generate dummy inputs because no PyTorch installation was found.")
260         elif framework == TensorType.TENSORFLOW:
261             if is_tf_available():
262                 return self._generate_extra_dummy_inputs_tf(dummy_inputs, label_batch_size, label_seq_length)
263             else:
264                 raise RuntimeError(f"Could not generate dummy inputs because no TensorFlow installation was found.")
265         else:
266             raise ValueError(
267                 f"Only two frameworks are supported for ONNX export: PyTorch or TensorFlow, but {framework} was provided."
268             )
269 
270 
271 class OnnxSeq2SeqConfigWithPastAndLoss(OnnxConfigWithPastAndLoss):
272     @property
273     def outputs(self) -> Mapping[str, Mapping[int, str]]:
274         common_outputs = self._onnx_config.outputs
275         extra_outputs = self._tasks_to_extra_outputs["default"]
276         common_outputs.update(extra_outputs)
277         for key in reversed(extra_outputs.keys()):
278             common_outputs.move_to_end(key, last=False)
279         return copy.deepcopy(common_outputs)
280
```
--- 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/onnx/configuration.py b/optimum/onnx/configuration.py
--- a/optimum/onnx/configuration.py
+++ b/optimum/onnx/configuration.py
@@ -168,15 +168,15 @@
         """
         # Generate dummy labels
         dummy_inputs = super().generate_dummy_inputs(
-            preprocessor,
-            batch_size,
-            seq_length,
-            is_pair,
-            framework,
-            num_channels,
-            image_width,
-            image_height,
-            tokenizer,
+            preprocessor=preprocessor,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            is_pair=is_pair,
+            framework=framework,
+            num_channels=num_channels,
+            image_width=image_width,
+            image_height=image_height,
+            tokenizer=tokenizer,
         )
         label_batch_size = compute_effective_axis_dimension(
             batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0
@@ -240,10 +240,10 @@
 
         dummy_inputs = self._onnx_config.generate_dummy_inputs(
             tokenizer,
-            batch_size,
-            seq_length,
-            is_pair,
-            framework,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            is_pair=is_pair,
+            framework=framework,
         )
         label_batch_size = compute_effective_axis_dimension(
             batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0

{"golden_diff": "diff --git a/optimum/onnx/configuration.py b/optimum/onnx/configuration.py\n--- a/optimum/onnx/configuration.py\n+++ b/optimum/onnx/configuration.py\n@@ -168,15 +168,15 @@\n         \"\"\"\n         # Generate dummy labels\n         dummy_inputs = super().generate_dummy_inputs(\n-            preprocessor,\n-            batch_size,\n-            seq_length,\n-            is_pair,\n-            framework,\n-            num_channels,\n-            image_width,\n-            image_height,\n-            tokenizer,\n+            preprocessor=preprocessor,\n+            batch_size=batch_size,\n+            seq_length=seq_length,\n+            is_pair=is_pair,\n+            framework=framework,\n+            num_channels=num_channels,\n+            image_width=image_width,\n+            image_height=image_height,\n+            tokenizer=tokenizer,\n         )\n         label_batch_size = compute_effective_axis_dimension(\n             batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0\n@@ -240,10 +240,10 @@\n \n         dummy_inputs = self._onnx_config.generate_dummy_inputs(\n             tokenizer,\n-            batch_size,\n-            seq_length,\n-            is_pair,\n-            framework,\n+            batch_size=batch_size,\n+            seq_length=seq_length,\n+            is_pair=is_pair,\n+            framework=framework,\n         )\n         label_batch_size = compute_effective_axis_dimension(\n             batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0\n", "issue": "`*with_loss` wrappers failed for transformers 4.19.0\n# Problem\r\nAs new the new argument `num_choices` added in `transformers.onnx.config.generate_dummy_inputs`, the wrappers will fail to match correct arguments\r\n\r\n### Error Message\r\n------------\r\n```\r\nTraceback (most recent call last):\r\n  File \"test_onnxruntime_train.py\", line 116, in test_ort_trainer\r\n    ort_eval_metrics = trainer.evaluate(inference_with_ort=inference_with_ort)\r\n  File \"/workspace/optimum/onnxruntime/trainer.py\", line 631, in evaluate\r\n    output = eval_loop(\r\n  File \"/workspace/optimum/onnxruntime/trainer.py\", line 767, in evaluation_loop_ort\r\n    self._export(onnx_model_path, with_loss=with_loss)\r\n  File \"/workspace/optimum/onnxruntime/trainer.py\", line 1230, in _export\r\n    _ = export(preprocessor=self.tokenizer, model=model, config=onnx_config, opset=opset, output=model_path)\r\n  File \"/usr/local/lib/python3.8/dist-packages/transformers/onnx/convert.py\", line 313, in export\r\n    return export_pytorch(preprocessor, model, config, opset, output, tokenizer=tokenizer)\r\n  File \"/usr/local/lib/python3.8/dist-packages/transformers/onnx/convert.py\", line 138, in export_pytorch\r\n    model_inputs = config.generate_dummy_inputs(preprocessor, framework=TensorType.PYTORCH)\r\n  File \"/workspace/optimum/onnx/configuration.py\", line 170, in generate_dummy_inputs\r\n    dummy_inputs = super().generate_dummy_inputs(\r\n  File \"/usr/local/lib/python3.8/dist-packages/transformers/onnx/config.py\", line 308, in generate_dummy_inputs\r\n    token_to_add = preprocessor.num_special_tokens_to_add(is_pair)\r\n  File \"/usr/local/lib/python3.8/dist-packages/transformers/tokenization_utils_fast.py\", line 289, in num_special_tokens_to_add\r\n    return self._tokenizer.num_special_tokens_to_add(pair)\r\nTypeError: Can't convert <TensorType.PYTORCH: 'pt'> to PyBool\r\n```\r\n\r\n# Solution\r\nAdd keys for arguments\n", "before_files": [{"content": "#  Copyright 2022 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\nimport copy\nfrom abc import ABC\nfrom collections import OrderedDict\nfrom typing import TYPE_CHECKING, Any, Mapping, Optional, Union\n\nfrom transformers.file_utils import TensorType, is_tf_available, is_torch_available\nfrom transformers.onnx.utils import compute_effective_axis_dimension\nfrom transformers.utils import logging\n\n\nif TYPE_CHECKING:\n    from transformers.configuration_utils import PretrainedConfig\n    from transformers.feature_extraction_utils import FeatureExtractionMixin\n    from transformers.tokenization_utils_base import PreTrainedTokenizerBase\n\nfrom transformers.onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast\n\n\nlogger = logging.get_logger(__name__)\n\n\nclass OnnxConfigWithLoss(OnnxConfig, ABC):\n    \"\"\"\n    Wrapper for the children classes of `transformers.onnx.OnnxConfig` to export the model through the ONNX format with loss in outputs.\n    \"\"\"\n\n    _tasks_to_extra_inputs = {\n        \"default\": OrderedDict({\"labels\": {0: \"batch\"}}),\n        \"masked-lm\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"causal-lm\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"seq2seq-lm\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"sequence-classification\": OrderedDict({\"labels\": {0: \"batch\"}}),\n        \"token-classification\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"multiple-choice\": OrderedDict({\"labels\": {0: \"batch\"}}),\n        \"question-answering\": OrderedDict(\n            {\n                \"start_positions\": {0: \"batch\"},\n                \"end_positions\": {0: \"batch\"},\n            }\n        ),\n        \"image-classification\": OrderedDict({\"labels\": {0: \"batch\"}}),\n    }\n    _tasks_to_extra_outputs = {\n        \"default\": OrderedDict({\"loss\": {}}),\n    }\n\n    def __init__(self, config: OnnxConfig):\n        self.__dict__ = copy.deepcopy(config.__dict__)\n        self._onnx_config = config\n        if self.task not in self._tasks_to_extra_inputs:\n            raise ValueError(\n                f\"{self.task} is not a supported task, supported tasks: {self._tasks_to_extra_inputs.keys()}\"\n            )\n\n    @classmethod\n    def from_model_config(cls, config: OnnxConfig) -> \"OnnxConfigWithLoss\":\n        \"\"\"\n        Instantiate a OnnxConfigWithLoss for a specific model\n        Args:\n            config: The model's configuration to use when exporting to ONNX\n        Returns:\n            OnnxConfigWithLoss for this model\n        \"\"\"\n        return cls(config)\n\n    @property\n    def inputs(self) -> Mapping[str, Mapping[int, str]]:\n        \"\"\"\n        Mapping containing the axis definition of the input tensors(including labels) to provide to the model\n        Returns:\n            For each input: its name associated to the axes symbolic name and the axis position within the tensor\n        \"\"\"\n        inputs = self._onnx_config.inputs\n        inputs.update(self._tasks_to_extra_inputs[self.task])\n        return inputs\n\n    @property\n    def outputs(self) -> Mapping[str, Mapping[int, str]]:\n        \"\"\"\n        Mapping containing the axis definition of the output tensors(including loss) to provide to the model\n        Returns:\n            For each output: its name associated to the axes symbolic name and the axis position within the tensor\n        \"\"\"\n        common_outputs = self._onnx_config.outputs\n        extra_outputs = self._tasks_to_extra_outputs[\"default\"]\n        common_outputs.update(extra_outputs)\n        for key in reversed(extra_outputs.keys()):\n            common_outputs.move_to_end(key, last=False)\n        return copy.deepcopy(common_outputs)\n\n    def _generate_extra_dummy_inputs_pt(\n        self,\n        dummy_inputs,\n        batch_size,\n        seq_length,\n    ) -> Mapping[str, Any]:\n        import torch\n\n        for label, input in self._tasks_to_extra_inputs[self.task].items():\n            if \"sequence\" in input.values():\n                dummy_inputs[label] = torch.zeros(batch_size, seq_length, dtype=torch.long)\n            else:\n                dummy_inputs[label] = torch.zeros(batch_size, dtype=torch.long)\n        return dummy_inputs\n\n    def _generate_extra_dummy_inputs_tf(\n        self,\n        dummy_inputs,\n        batch_size,\n        seq_length,\n    ) -> Mapping[str, Any]:\n        import tensorflow as tf\n\n        for label, input in self._tasks_to_extra_inputs[self.task].items():\n            if \"sequence\" in input.values():\n                dummy_inputs[label] = tf.zeros(batch_size, seq_length, dtype=tf.int64)\n            else:\n                dummy_inputs[label] = tf.zeros(batch_size, dtype=tf.int64)\n        return dummy_inputs\n\n    def generate_dummy_inputs(\n        self,\n        preprocessor: Union[\"PreTrainedTokenizerBase\", \"FeatureExtractionMixin\"],\n        batch_size: int = -1,\n        seq_length: int = -1,\n        is_pair: bool = False,\n        framework: Optional[TensorType] = None,\n        num_channels: int = 3,\n        image_width: int = 40,\n        image_height: int = 40,\n        tokenizer: \"PreTrainedTokenizerBase\" = None,\n    ) -> Mapping[str, Any]:\n        \"\"\"\n        Generate inputs to provide to the ONNX exporter for the specific framework\n        Args:\n            preprocessor: ([`PreTrainedTokenizerBase`] or [`FeatureExtractionMixin`]):\n                The preprocessor associated with this model configuration.\n            batch_size (`int`, *optional*, defaults to -1):\n                The batch size to export the model for (-1 means dynamic axis).\n            seq_length (`int`, *optional*, defaults to -1):\n                The sequence length to export the model for (-1 means dynamic axis).\n            is_pair (`bool`, *optional*, defaults to `False`):\n                Indicate if the input is a pair (sentence 1, sentence 2)\n            framework (`TensorType`, *optional*, defaults to `None`):\n                The framework (PyTorch or TensorFlow) that the tokenizer will generate tensors for.\n            num_channels (`int`, *optional*, defaults to 3):\n                The number of channels of the generated images.\n            image_width (`int`, *optional*, defaults to 40):\n                The width of the generated images.\n            image_height (`int`, *optional*, defaults to 40):\n                The height of the generated images.\n        Returns:\n            Mapping[str, Tensor] holding the kwargs to provide to the model's forward function\n        \"\"\"\n        # Generate dummy labels\n        dummy_inputs = super().generate_dummy_inputs(\n            preprocessor,\n            batch_size,\n            seq_length,\n            is_pair,\n            framework,\n            num_channels,\n            image_width,\n            image_height,\n            tokenizer,\n        )\n        label_batch_size = compute_effective_axis_dimension(\n            batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0\n        )\n        label_seq_length = compute_effective_axis_dimension(\n            seq_length, fixed_dimension=self.default_fixed_sequence, num_token_to_add=0\n        )\n\n        if framework == TensorType.PYTORCH:\n            if is_torch_available():\n                return self._generate_extra_dummy_inputs_pt(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no PyTorch installation was found.\")\n        elif framework == TensorType.TENSORFLOW:\n            if is_tf_available():\n                return self._generate_extra_dummy_inputs_tf(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no TensorFlow installation was found.\")\n        else:\n            raise ValueError(\n                f\"Only two frameworks are supported for ONNX export: PyTorch or TensorFlow, but {framework} was provided.\"\n            )\n\n\nclass OnnxConfigWithPastAndLoss(OnnxConfigWithLoss, ABC):\n    def __init__(\n        self,\n        config: OnnxConfigWithPast,\n        use_past: bool = False,\n    ):\n        super().__init__(config)\n        self.use_past = use_past\n\n    @classmethod\n    def with_past(cls, config: OnnxConfigWithPast) -> \"OnnxConfigWithPast\":\n        \"\"\"\n        Instantiate a OnnxConfigWithPastAndLoss with `use_past` attribute set to True\n        Args:\n            config: The underlying model's config to use when exporting to ONNX\n        Returns:\n            OnnxConfigWithPastAndLoss with `.use_past = True`\n        \"\"\"\n        return cls(config, use_past=True)\n\n    @property\n    def outputs(self) -> Mapping[str, Mapping[int, str]]:\n        common_outputs = super().outputs\n        if self.use_past:\n            self._onnx_config.fill_with_past_key_values_(common_outputs, direction=\"outputs\")\n\n        return common_outputs\n\n    def generate_dummy_inputs(\n        self,\n        tokenizer: \"PreTrainedTokenizerBase\",\n        batch_size: int = -1,\n        seq_length: int = -1,\n        is_pair: bool = False,\n        framework: Optional[TensorType] = None,\n    ) -> Mapping[str, Any]:\n\n        dummy_inputs = self._onnx_config.generate_dummy_inputs(\n            tokenizer,\n            batch_size,\n            seq_length,\n            is_pair,\n            framework,\n        )\n        label_batch_size = compute_effective_axis_dimension(\n            batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0\n        )\n        label_seq_length = compute_effective_axis_dimension(\n            seq_length, fixed_dimension=self.default_fixed_sequence, num_token_to_add=0\n        )\n\n        if framework == TensorType.PYTORCH:\n            if is_torch_available():\n                return self._generate_extra_dummy_inputs_pt(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no PyTorch installation was found.\")\n        elif framework == TensorType.TENSORFLOW:\n            if is_tf_available():\n                return self._generate_extra_dummy_inputs_tf(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no TensorFlow installation was found.\")\n        else:\n            raise ValueError(\n                f\"Only two frameworks are supported for ONNX export: PyTorch or TensorFlow, but {framework} was provided.\"\n            )\n\n\nclass OnnxSeq2SeqConfigWithPastAndLoss(OnnxConfigWithPastAndLoss):\n    @property\n    def outputs(self) -> Mapping[str, Mapping[int, str]]:\n        common_outputs = self._onnx_config.outputs\n        extra_outputs = self._tasks_to_extra_outputs[\"default\"]\n        common_outputs.update(extra_outputs)\n        for key in reversed(extra_outputs.keys()):\n            common_outputs.move_to_end(key, last=False)\n        return copy.deepcopy(common_outputs)\n", "path": "optimum/onnx/configuration.py"}], "after_files": [{"content": "#  Copyright 2022 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\nimport copy\nfrom abc import ABC\nfrom collections import OrderedDict\nfrom typing import TYPE_CHECKING, Any, Mapping, Optional, Union\n\nfrom transformers.file_utils import TensorType, is_tf_available, is_torch_available\nfrom transformers.onnx.utils import compute_effective_axis_dimension\nfrom transformers.utils import logging\n\n\nif TYPE_CHECKING:\n    from transformers.configuration_utils import PretrainedConfig\n    from transformers.feature_extraction_utils import FeatureExtractionMixin\n    from transformers.tokenization_utils_base import PreTrainedTokenizerBase\n\nfrom transformers.onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast\n\n\nlogger = logging.get_logger(__name__)\n\n\nclass OnnxConfigWithLoss(OnnxConfig, ABC):\n    \"\"\"\n    Wrapper for the children classes of `transformers.onnx.OnnxConfig` to export the model through the ONNX format with loss in outputs.\n    \"\"\"\n\n    _tasks_to_extra_inputs = {\n        \"default\": OrderedDict({\"labels\": {0: \"batch\"}}),\n        \"masked-lm\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"causal-lm\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"seq2seq-lm\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"sequence-classification\": OrderedDict({\"labels\": {0: \"batch\"}}),\n        \"token-classification\": OrderedDict({\"labels\": {0: \"batch\", 1: \"sequence\"}}),\n        \"multiple-choice\": OrderedDict({\"labels\": {0: \"batch\"}}),\n        \"question-answering\": OrderedDict(\n            {\n                \"start_positions\": {0: \"batch\"},\n                \"end_positions\": {0: \"batch\"},\n            }\n        ),\n        \"image-classification\": OrderedDict({\"labels\": {0: \"batch\"}}),\n    }\n    _tasks_to_extra_outputs = {\n        \"default\": OrderedDict({\"loss\": {}}),\n    }\n\n    def __init__(self, config: OnnxConfig):\n        self.__dict__ = copy.deepcopy(config.__dict__)\n        self._onnx_config = config\n        if self.task not in self._tasks_to_extra_inputs:\n            raise ValueError(\n                f\"{self.task} is not a supported task, supported tasks: {self._tasks_to_extra_inputs.keys()}\"\n            )\n\n    @classmethod\n    def from_model_config(cls, config: OnnxConfig) -> \"OnnxConfigWithLoss\":\n        \"\"\"\n        Instantiate a OnnxConfigWithLoss for a specific model\n        Args:\n            config: The model's configuration to use when exporting to ONNX\n        Returns:\n            OnnxConfigWithLoss for this model\n        \"\"\"\n        return cls(config)\n\n    @property\n    def inputs(self) -> Mapping[str, Mapping[int, str]]:\n        \"\"\"\n        Mapping containing the axis definition of the input tensors(including labels) to provide to the model\n        Returns:\n            For each input: its name associated to the axes symbolic name and the axis position within the tensor\n        \"\"\"\n        inputs = self._onnx_config.inputs\n        inputs.update(self._tasks_to_extra_inputs[self.task])\n        return inputs\n\n    @property\n    def outputs(self) -> Mapping[str, Mapping[int, str]]:\n        \"\"\"\n        Mapping containing the axis definition of the output tensors(including loss) to provide to the model\n        Returns:\n            For each output: its name associated to the axes symbolic name and the axis position within the tensor\n        \"\"\"\n        common_outputs = self._onnx_config.outputs\n        extra_outputs = self._tasks_to_extra_outputs[\"default\"]\n        common_outputs.update(extra_outputs)\n        for key in reversed(extra_outputs.keys()):\n            common_outputs.move_to_end(key, last=False)\n        return copy.deepcopy(common_outputs)\n\n    def _generate_extra_dummy_inputs_pt(\n        self,\n        dummy_inputs,\n        batch_size,\n        seq_length,\n    ) -> Mapping[str, Any]:\n        import torch\n\n        for label, input in self._tasks_to_extra_inputs[self.task].items():\n            if \"sequence\" in input.values():\n                dummy_inputs[label] = torch.zeros(batch_size, seq_length, dtype=torch.long)\n            else:\n                dummy_inputs[label] = torch.zeros(batch_size, dtype=torch.long)\n        return dummy_inputs\n\n    def _generate_extra_dummy_inputs_tf(\n        self,\n        dummy_inputs,\n        batch_size,\n        seq_length,\n    ) -> Mapping[str, Any]:\n        import tensorflow as tf\n\n        for label, input in self._tasks_to_extra_inputs[self.task].items():\n            if \"sequence\" in input.values():\n                dummy_inputs[label] = tf.zeros(batch_size, seq_length, dtype=tf.int64)\n            else:\n                dummy_inputs[label] = tf.zeros(batch_size, dtype=tf.int64)\n        return dummy_inputs\n\n    def generate_dummy_inputs(\n        self,\n        preprocessor: Union[\"PreTrainedTokenizerBase\", \"FeatureExtractionMixin\"],\n        batch_size: int = -1,\n        seq_length: int = -1,\n        is_pair: bool = False,\n        framework: Optional[TensorType] = None,\n        num_channels: int = 3,\n        image_width: int = 40,\n        image_height: int = 40,\n        tokenizer: \"PreTrainedTokenizerBase\" = None,\n    ) -> Mapping[str, Any]:\n        \"\"\"\n        Generate inputs to provide to the ONNX exporter for the specific framework\n        Args:\n            preprocessor: ([`PreTrainedTokenizerBase`] or [`FeatureExtractionMixin`]):\n                The preprocessor associated with this model configuration.\n            batch_size (`int`, *optional*, defaults to -1):\n                The batch size to export the model for (-1 means dynamic axis).\n            seq_length (`int`, *optional*, defaults to -1):\n                The sequence length to export the model for (-1 means dynamic axis).\n            is_pair (`bool`, *optional*, defaults to `False`):\n                Indicate if the input is a pair (sentence 1, sentence 2)\n            framework (`TensorType`, *optional*, defaults to `None`):\n                The framework (PyTorch or TensorFlow) that the tokenizer will generate tensors for.\n            num_channels (`int`, *optional*, defaults to 3):\n                The number of channels of the generated images.\n            image_width (`int`, *optional*, defaults to 40):\n                The width of the generated images.\n            image_height (`int`, *optional*, defaults to 40):\n                The height of the generated images.\n        Returns:\n            Mapping[str, Tensor] holding the kwargs to provide to the model's forward function\n        \"\"\"\n        # Generate dummy labels\n        dummy_inputs = super().generate_dummy_inputs(\n            preprocessor=preprocessor,\n            batch_size=batch_size,\n            seq_length=seq_length,\n            is_pair=is_pair,\n            framework=framework,\n            num_channels=num_channels,\n            image_width=image_width,\n            image_height=image_height,\n            tokenizer=tokenizer,\n        )\n        label_batch_size = compute_effective_axis_dimension(\n            batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0\n        )\n        label_seq_length = compute_effective_axis_dimension(\n            seq_length, fixed_dimension=self.default_fixed_sequence, num_token_to_add=0\n        )\n\n        if framework == TensorType.PYTORCH:\n            if is_torch_available():\n                return self._generate_extra_dummy_inputs_pt(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no PyTorch installation was found.\")\n        elif framework == TensorType.TENSORFLOW:\n            if is_tf_available():\n                return self._generate_extra_dummy_inputs_tf(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no TensorFlow installation was found.\")\n        else:\n            raise ValueError(\n                f\"Only two frameworks are supported for ONNX export: PyTorch or TensorFlow, but {framework} was provided.\"\n            )\n\n\nclass OnnxConfigWithPastAndLoss(OnnxConfigWithLoss, ABC):\n    def __init__(\n        self,\n        config: OnnxConfigWithPast,\n        use_past: bool = False,\n    ):\n        super().__init__(config)\n        self.use_past = use_past\n\n    @classmethod\n    def with_past(cls, config: OnnxConfigWithPast) -> \"OnnxConfigWithPast\":\n        \"\"\"\n        Instantiate a OnnxConfigWithPastAndLoss with `use_past` attribute set to True\n        Args:\n            config: The underlying model's config to use when exporting to ONNX\n        Returns:\n            OnnxConfigWithPastAndLoss with `.use_past = True`\n        \"\"\"\n        return cls(config, use_past=True)\n\n    @property\n    def outputs(self) -> Mapping[str, Mapping[int, str]]:\n        common_outputs = super().outputs\n        if self.use_past:\n            self._onnx_config.fill_with_past_key_values_(common_outputs, direction=\"outputs\")\n\n        return common_outputs\n\n    def generate_dummy_inputs(\n        self,\n        tokenizer: \"PreTrainedTokenizerBase\",\n        batch_size: int = -1,\n        seq_length: int = -1,\n        is_pair: bool = False,\n        framework: Optional[TensorType] = None,\n    ) -> Mapping[str, Any]:\n\n        dummy_inputs = self._onnx_config.generate_dummy_inputs(\n            tokenizer,\n            batch_size=batch_size,\n            seq_length=seq_length,\n            is_pair=is_pair,\n            framework=framework,\n        )\n        label_batch_size = compute_effective_axis_dimension(\n            batch_size, fixed_dimension=self.default_fixed_batch, num_token_to_add=0\n        )\n        label_seq_length = compute_effective_axis_dimension(\n            seq_length, fixed_dimension=self.default_fixed_sequence, num_token_to_add=0\n        )\n\n        if framework == TensorType.PYTORCH:\n            if is_torch_available():\n                return self._generate_extra_dummy_inputs_pt(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no PyTorch installation was found.\")\n        elif framework == TensorType.TENSORFLOW:\n            if is_tf_available():\n                return self._generate_extra_dummy_inputs_tf(dummy_inputs, label_batch_size, label_seq_length)\n            else:\n                raise RuntimeError(f\"Could not generate dummy inputs because no TensorFlow installation was found.\")\n        else:\n            raise ValueError(\n                f\"Only two frameworks are supported for ONNX export: PyTorch or TensorFlow, but {framework} was provided.\"\n            )\n\n\nclass OnnxSeq2SeqConfigWithPastAndLoss(OnnxConfigWithPastAndLoss):\n    @property\n    def outputs(self) -> Mapping[str, Mapping[int, str]]:\n        common_outputs = self._onnx_config.outputs\n        extra_outputs = self._tasks_to_extra_outputs[\"default\"]\n        common_outputs.update(extra_outputs)\n        for key in reversed(extra_outputs.keys()):\n            common_outputs.move_to_end(key, last=False)\n        return copy.deepcopy(common_outputs)\n", "path": "optimum/onnx/configuration.py"}]}