in optimum/exporters/onnx/model_configs.py [0:0]
def inputs(self) -> Dict[str, Dict[int, str]]:
common_inputs = {}
if self._behavior is not ConfigBehavior.DECODER:
common_inputs["input_features"] = {0: "batch_size", 1: "feature_size", 2: "encoder_sequence_length"}
common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"}
if self._behavior is not ConfigBehavior.ENCODER:
if self.use_past_in_inputs:
common_inputs["decoder_input_ids"] = {0: "batch_size"}
else:
common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"}
if self.use_past_in_inputs:
self.add_past_key_values(common_inputs, direction="inputs")
if self._behavior is ConfigBehavior.DECODER:
common_inputs["encoder_outputs"] = {
0: "batch_size",
1: f"encoder_sequence_length / {(2 * self._config.num_conv_layers)}",
}
return common_inputs