

from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


lowerCamelCase_    =					logging.get_logger(__name__)

lowerCamelCase_    =					{
    '''facebook/mbart-large-cc25''': '''https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json''',
    # See all MBART models at https://huggingface.co/models?filter=mbart
}







class  __lowerCamelCase				(	__snake_case	):
      lowerCamelCase_	:      Union[str, Any]          =					'mbart'
      lowerCamelCase_	:      int          =					['past_key_values']
      lowerCamelCase_	:      List[Any]          =					{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}







      def __init__(     self							,  lowerCamelCase=50265							,  lowerCamelCase=1024							,  lowerCamelCase=12							,  lowerCamelCase=4096							,  lowerCamelCase=16							,  lowerCamelCase=12							,  lowerCamelCase=4096							,  lowerCamelCase=16							,  lowerCamelCase=0.0							,  lowerCamelCase=0.0							,  lowerCamelCase=True							,  lowerCamelCase=True							,  lowerCamelCase="gelu"							,  lowerCamelCase=1024							,  lowerCamelCase=0.1							,  lowerCamelCase=0.0							,  lowerCamelCase=0.0							,  lowerCamelCase=0.02							,  lowerCamelCase=0.0							,  lowerCamelCase=False							,  lowerCamelCase=1							,  lowerCamelCase=0							,  lowerCamelCase=2							,  lowerCamelCase=2							,  **lowerCamelCase							,  )						->       Tuple:
        snake_case_       						=      vocab_size
        snake_case_       						=      max_position_embeddings
        snake_case_       						=      d_model
        snake_case_       						=      encoder_ffn_dim
        snake_case_       						=      encoder_layers
        snake_case_       						=      encoder_attention_heads
        snake_case_       						=      decoder_ffn_dim
        snake_case_       						=      decoder_layers
        snake_case_       						=      decoder_attention_heads
        snake_case_       						=      dropout
        snake_case_       						=      attention_dropout
        snake_case_       						=      activation_dropout
        snake_case_       						=      activation_function
        snake_case_       						=      init_std
        snake_case_       						=      encoder_layerdrop
        snake_case_       						=      decoder_layerdrop
        snake_case_       						=      classifier_dropout
        snake_case_       						=      use_cache
        snake_case_       						=      encoder_layers
        snake_case_       						=      scale_embedding  # scale factor will be sqrt(d_model) if True
        super().__init__(
            pad_token_id=lowerCamelCase							,  bos_token_id=lowerCamelCase							,  eos_token_id=lowerCamelCase							,  is_encoder_decoder=lowerCamelCase							,  forced_eos_token_id=lowerCamelCase							,  **lowerCamelCase							,  )







class  __lowerCamelCase				(	__snake_case	):







      @property
      def 	lowerCAmelCase_    (     self   )						->       Mapping[str, Mapping[int, str]]:
        if self.task in ["default", "seq2seq-lm"]:
          snake_case_       						=      OrderedDict(
              [
                  ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                  ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
              ]   )

          if self.use_past:
            snake_case_       						=      {0: """batch"""}
            snake_case_       						=      {0: """batch""", 1: """past_decoder_sequence + sequence"""}
          else:
            snake_case_       						=      {0: """batch""", 1: """decoder_sequence"""}
            snake_case_       						=      {0: """batch""", 1: """decoder_sequence"""}

          if self.use_past:
            self.fill_with_past_key_values_(lowerCamelCase							,  direction="""inputs"""   )
        elif self.task == "causal-lm":
          # TODO: figure this case out.
          snake_case_       						=      OrderedDict(
              [
                  ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                  ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
              ]   )
          if self.use_past:
            snake_case_						,				snake_case_       						=      self.num_layers
            for i in range(lowerCamelCase   ):
              snake_case_       						=      {0: """batch""", 2: """past_sequence + sequence"""}
              snake_case_       						=      {0: """batch""", 2: """past_sequence + sequence"""}
        else:
          snake_case_       						=      OrderedDict(
              [
                  ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                  ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                  ("""decoder_input_ids""", {0: """batch""", 1: """decoder_sequence"""}),
                  ("""decoder_attention_mask""", {0: """batch""", 1: """decoder_sequence"""}),
              ]   )

        return common_inputs







      @property
      def 	lowerCAmelCase_    (     self   )						->       Mapping[str, Mapping[int, str]]:
        if self.task in ["default", "seq2seq-lm"]:
          snake_case_       						=      super().outputs
        else:
          snake_case_       						=      super(lowerCamelCase							,  self   ).outputs
          if self.use_past:
            snake_case_						,				snake_case_       						=      self.num_layers
            for i in range(lowerCamelCase   ):
              snake_case_       						=      {0: """batch""", 2: """past_sequence + sequence"""}
              snake_case_       						=      {0: """batch""", 2: """past_sequence + sequence"""}
        return common_outputs







      def 	lowerCAmelCase_    (     self							,  lowerCamelCase							,  lowerCamelCase = -1							,  lowerCamelCase = -1							,  lowerCamelCase = False							,  lowerCamelCase = None							,  )						->       Mapping[str, Any]:
        snake_case_       						=      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase   )

        # Generate decoder inputs
        snake_case_       						=      seq_length if not self.use_past else 1
        snake_case_       						=      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase   )
        snake_case_       						=      {f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
        snake_case_       						=      dict(**lowerCamelCase							,  **lowerCamelCase   )

        if self.use_past:
          if not is_torch_available():
            raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""   )
          else:
            import torch
          snake_case_						,				snake_case_       						=      common_inputs["""input_ids"""].shape
          snake_case_       						=      common_inputs["""decoder_input_ids"""].shape[1]
          snake_case_						,				snake_case_       						=      self.num_attention_heads
          snake_case_       						=      (
              batch,
              num_encoder_attention_heads,
              encoder_seq_length,
              self._config.hidden_size // num_encoder_attention_heads,
          )
          snake_case_       						=      decoder_seq_length + 3
          snake_case_       						=      (
              batch,
              num_decoder_attention_heads,
              decoder_past_length,
              self._config.hidden_size // num_decoder_attention_heads,
          )

          snake_case_       						=      torch.cat(
              [common_inputs["""decoder_attention_mask"""], torch.ones(lowerCamelCase							,  lowerCamelCase   )]							,  dim=1   )

          snake_case_       						=      []
          # If the number of encoder and decoder layers are present in the model configuration, both are considered
          snake_case_						,				snake_case_       						=      self.num_layers
          snake_case_       						=      min(lowerCamelCase							,  lowerCamelCase   )
          snake_case_       						=      max(lowerCamelCase							,  lowerCamelCase   ) - min_num_layers
          snake_case_       						=      """encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

          for _ in range(lowerCamelCase   ):
            common_inputs["past_key_values"].append(
                (
                    torch.zeros(lowerCamelCase   ),
                    torch.zeros(lowerCamelCase   ),
                    torch.zeros(lowerCamelCase   ),
                    torch.zeros(lowerCamelCase   ),
                )   )
          # TODO: test this.
          snake_case_       						=      encoder_shape if remaining_side_name == """encoder""" else decoder_shape
          for _ in range(lowerCamelCase							,  lowerCamelCase   ):
            common_inputs["past_key_values"].append((torch.zeros(lowerCamelCase   ), torch.zeros(lowerCamelCase   ))   )
        return common_inputs







      def 	lowerCAmelCase_    (     self							,  lowerCamelCase							,  lowerCamelCase = -1							,  lowerCamelCase = -1							,  lowerCamelCase = False							,  lowerCamelCase = None							,  )						->       Mapping[str, Any]:
        snake_case_       						=      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase   )

        if self.use_past:
          if not is_torch_available():
            raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""   )
          else:
            import torch
          snake_case_						,				snake_case_       						=      common_inputs["""input_ids"""].shape
          # Not using the same length for past_key_values
          snake_case_       						=      seqlen + 2
          snake_case_						,				snake_case_       						=      self.num_layers
          snake_case_						,				snake_case_       						=      self.num_attention_heads
          snake_case_       						=      (
              batch,
              num_encoder_attention_heads,
              past_key_values_length,
              self._config.hidden_size // num_encoder_attention_heads,
          )

          snake_case_       						=      common_inputs["""attention_mask"""].dtype
          snake_case_       						=      torch.cat(
              [common_inputs["""attention_mask"""], torch.ones(lowerCamelCase							,  lowerCamelCase							,  dtype=lowerCamelCase   )]							,  dim=1   )
          snake_case_       						=      [
              (torch.zeros(lowerCamelCase   ), torch.zeros(lowerCamelCase   )) for _ in range(lowerCamelCase   )
          ]
        return common_inputs







      def 	lowerCAmelCase_    (     self							,  lowerCamelCase							,  lowerCamelCase = -1							,  lowerCamelCase = -1							,  lowerCamelCase = False							,  lowerCamelCase = None							,  )						->       Mapping[str, Any]:
        # Copied from OnnxConfig.generate_dummy_inputs
        # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
        # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
        snake_case_       						=      compute_effective_axis_dimension(
            lowerCamelCase							,  fixed_dimension=OnnxConfig.default_fixed_batch							,  num_token_to_add=0   )

        # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
        snake_case_       						=      tokenizer.num_special_tokens_to_add(lowerCamelCase   )
        snake_case_       						=      compute_effective_axis_dimension(
            lowerCamelCase							,  fixed_dimension=OnnxConfig.default_fixed_sequence							,  num_token_to_add=lowerCamelCase   )

        # Generate dummy inputs according to compute batch and sequence
        snake_case_       						=      [""" """.join([tokenizer.unk_token]   ) * seq_length] * batch_size
        snake_case_       						=      dict(tokenizer(lowerCamelCase							,  return_tensors=lowerCamelCase   )   )
        return common_inputs







      def 	lowerCAmelCase_    (     self							,  lowerCamelCase							,  lowerCamelCase = -1							,  lowerCamelCase = -1							,  lowerCamelCase = False							,  lowerCamelCase = None							,  )						->       Mapping[str, Any]:
        if self.task in ["default", "seq2seq-lm"]:
          snake_case_       						=      self._generate_dummy_inputs_for_default_and_seqaseq_lm(
              lowerCamelCase							,  batch_size=lowerCamelCase							,  seq_length=lowerCamelCase							,  is_pair=lowerCamelCase							,  framework=lowerCamelCase   )

        elif self.task == "causal-lm":
          snake_case_       						=      self._generate_dummy_inputs_for_causal_lm(
              lowerCamelCase							,  batch_size=lowerCamelCase							,  seq_length=lowerCamelCase							,  is_pair=lowerCamelCase							,  framework=lowerCamelCase   )
        else:
          snake_case_       						=      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              lowerCamelCase							,  batch_size=lowerCamelCase							,  seq_length=lowerCamelCase							,  is_pair=lowerCamelCase							,  framework=lowerCamelCase   )

        return common_inputs







      def 	lowerCAmelCase_    (     self							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase   )						->       Dict:
        if self.task in ["default", "seq2seq-lm"]:
          snake_case_       						=      super()._flatten_past_key_values_(lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase   )
        else:
          snake_case_       						=      super(lowerCamelCase							,  self   )._flatten_past_key_values_(
              lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase							,  lowerCamelCase   )






from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


UpperCAmelCase__  :	str     =logging.get_logger(__name__)

UpperCAmelCase__  :	Optional[int]     ={
    '''gpt2''': '''https://huggingface.co/gpt2/resolve/main/config.json''',
    '''gpt2-medium''': '''https://huggingface.co/gpt2-medium/resolve/main/config.json''',
    '''gpt2-large''': '''https://huggingface.co/gpt2-large/resolve/main/config.json''',
    '''gpt2-xl''': '''https://huggingface.co/gpt2-xl/resolve/main/config.json''',
    '''distilgpt2''': '''https://huggingface.co/distilgpt2/resolve/main/config.json''',
}

class 					__A   (  a			):
				__A            = """gpt2"""
				__A            = ["""past_key_values"""]
				__A            = {
				    """hidden_size""": """n_embd""",
				    """max_position_embeddings""": """n_positions""",
				    """num_attention_heads""": """n_head""",
				    """num_hidden_layers""": """n_layer""",
				}


				def __init__(						self      ,       UpperCAmelCase_=50257      ,       UpperCAmelCase_=1024      ,       UpperCAmelCase_=768      ,       UpperCAmelCase_=12      ,       UpperCAmelCase_=12      ,       UpperCAmelCase_=None      ,       UpperCAmelCase_="gelu_new"      ,       UpperCAmelCase_=0.1      ,       UpperCAmelCase_=0.1      ,       UpperCAmelCase_=0.1      ,       UpperCAmelCase_=1E-5      ,       UpperCAmelCase_=0.0_2      ,       UpperCAmelCase_="cls_index"      ,       UpperCAmelCase_=True      ,       UpperCAmelCase_=None      ,       UpperCAmelCase_=True      ,       UpperCAmelCase_=0.1      ,       UpperCAmelCase_=True      ,       UpperCAmelCase_=True      ,       UpperCAmelCase_=50256      ,       UpperCAmelCase_=50256      ,       UpperCAmelCase_=False      ,       UpperCAmelCase_=False      ,       **UpperCAmelCase_      ,       ):
								lowerCamelCase		       =vocab_size
								lowerCamelCase		       =n_positions
								lowerCamelCase		       =n_embd
								lowerCamelCase		       =n_layer
								lowerCamelCase		       =n_head
								lowerCamelCase		       =n_inner
								lowerCamelCase		       =activation_function
								lowerCamelCase		       =resid_pdrop
								lowerCamelCase		       =embd_pdrop
								lowerCamelCase		       =attn_pdrop
								lowerCamelCase		       =layer_norm_epsilon
								lowerCamelCase		       =initializer_range
								lowerCamelCase		       =summary_type
								lowerCamelCase		       =summary_use_proj
								lowerCamelCase		       =summary_activation
								lowerCamelCase		       =summary_first_dropout
								lowerCamelCase		       =summary_proj_to_labels
								lowerCamelCase		       =scale_attn_weights
								lowerCamelCase		       =use_cache
								lowerCamelCase		       =scale_attn_by_inverse_layer_idx
								lowerCamelCase		       =reorder_and_upcast_attn

								lowerCamelCase		       =bos_token_id
								lowerCamelCase		       =eos_token_id

								super().__init__(bos_token_id=UpperCAmelCase_      ,       eos_token_id=UpperCAmelCase_      ,       **UpperCAmelCase_							)




class 					__A   (  a			):


				def __init__(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = "default"      ,       UpperCAmelCase_ = None      ,       UpperCAmelCase_ = False      ,       ):
								super().__init__(UpperCAmelCase_      ,       task=UpperCAmelCase_      ,       patching_specs=UpperCAmelCase_      ,       use_past=UpperCAmelCase_							)
								if not getattr(self._config      ,       """pad_token_id"""      ,       UpperCAmelCase_							):
												# TODO: how to do that better?
												lowerCamelCase		       =0


				@property
				def _snake_case				(						self							):
								lowerCamelCase		       =OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}							)
								if self.use_past:
												self.fill_with_past_key_values_(UpperCAmelCase_      ,       direction="""inputs"""							)
												lowerCamelCase		       ={0: """batch""", 1: """past_sequence + sequence"""}
								else:
												lowerCamelCase		       ={0: """batch""", 1: """sequence"""}

								return common_inputs


				@property
				def _snake_case				(						self							):
								return self._config.n_layer


				@property
				def _snake_case				(						self							):
								return self._config.n_head


				def _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = -1      ,       UpperCAmelCase_ = -1      ,       UpperCAmelCase_ = False      ,       UpperCAmelCase_ = None      ,       ):
								lowerCamelCase		       =super(UpperCAmelCase_      ,       self							).generate_dummy_inputs(
								    UpperCAmelCase_      ,       batch_size=UpperCAmelCase_      ,       seq_length=UpperCAmelCase_      ,       is_pair=UpperCAmelCase_      ,       framework=UpperCAmelCase_							)

								# We need to order the input in the way they appears in the forward()
								lowerCamelCase		       =OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}							)

								# Need to add the past_keys
								if self.use_past:
												if not is_torch_available():
																raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""							)
												else:
																import torch

																lowerCamelCase       ,     lowerCamelCase		       =common_inputs["""input_ids"""].shape
																# Not using the same length for past_key_values
																lowerCamelCase		       =seqlen + 2
																lowerCamelCase		       =(
																    batch,
																    self.num_attention_heads,
																    past_key_values_length,
																    self._config.hidden_size // self.num_attention_heads,
																)
																lowerCamelCase		       =[
																    (torch.zeros(UpperCAmelCase_							), torch.zeros(UpperCAmelCase_							)) for _ in range(self.num_layers							)
																]

								lowerCamelCase		       =common_inputs["""attention_mask"""]
								if self.use_past:
												lowerCamelCase		       =ordered_inputs["""attention_mask"""].dtype
												lowerCamelCase		       =torch.cat(
												    [ordered_inputs["""attention_mask"""], torch.ones(UpperCAmelCase_      ,       UpperCAmelCase_      ,       dtype=UpperCAmelCase_							)]      ,       dim=1							)

								return ordered_inputs


				@property
				def _snake_case				(						self							):
								return 13





