


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


SCREAMING_SNAKE_CASE	:Dict  		=     logging.get_logger(__name__)

SCREAMING_SNAKE_CASE	:int  		=     {
    '''google/bigbird-pegasus-large-arxiv''': (
        '''https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json'''
    ),
    '''google/bigbird-pegasus-large-pubmed''': (
        '''https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json'''
    ),
    '''google/bigbird-pegasus-large-bigpatent''': (
        '''https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json'''
    ),
    # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
}


class     __lowerCAmelCase     ( a					):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		    =       'bigbird_pegasus'
  _SCREAMING_SNAKE_CASE		    =       ['past_key_values']
  _SCREAMING_SNAKE_CASE		    =       {
      'num_attention_heads': 'encoder_attention_heads',
      'hidden_size': 'd_model',
      'attention_probs_dropout_prob': 'attention_dropout',
  }



  def __init__(     self       :       str    ,				_lowerCAmelCase       :       Union[str, Any]=9_6_1_0_3    ,				_lowerCAmelCase       :       int=4_0_9_6    ,				_lowerCAmelCase       :       int=1_6    ,				_lowerCAmelCase       :       Optional[Any]=4_0_9_6    ,				_lowerCAmelCase       :       List[Any]=1_6    ,				_lowerCAmelCase       :       Optional[Any]=1_6    ,				_lowerCAmelCase       :       Tuple=4_0_9_6    ,				_lowerCAmelCase       :       List[Any]=1_6    ,				_lowerCAmelCase       :       str=0.0    ,				_lowerCAmelCase       :       Dict=0.0    ,				_lowerCAmelCase       :       List[str]=True    ,				_lowerCAmelCase       :       List[Any]=True    ,				_lowerCAmelCase       :       List[Any]="gelu_new"    ,				_lowerCAmelCase       :       int=1_0_2_4    ,				_lowerCAmelCase       :       Any=0.1    ,				_lowerCAmelCase       :       Union[str, Any]=0.0    ,				_lowerCAmelCase       :       str=0.0    ,				_lowerCAmelCase       :       Union[str, Any]=0.02    ,				_lowerCAmelCase       :       Any=2    ,				_lowerCAmelCase       :       List[str]=0.0    ,				_lowerCAmelCase       :       Any=True    ,				_lowerCAmelCase       :       List[str]=0    ,				_lowerCAmelCase       :       Dict=2    ,				_lowerCAmelCase       :       Tuple=1    ,				_lowerCAmelCase       :       Dict="block_sparse"    ,				_lowerCAmelCase       :       Optional[Any]=6_4    ,				_lowerCAmelCase       :       List[Any]=3    ,				_lowerCAmelCase       :       int=False    ,				**_lowerCAmelCase       :       List[Any]    ,				)   -> str:




     """simple docstring"""
     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

     # extra config
     snake_case_					    =					attention_type
     snake_case_					    =					block_size
     snake_case_					    =					num_random_blocks
     snake_case_					    =					use_bias

     super().__init__(
         pad_token_id=_lowerCAmelCase    ,				bos_token_id=_lowerCAmelCase    ,				eos_token_id=_lowerCAmelCase    ,				is_encoder_decoder=_lowerCAmelCase    ,				decoder_start_token_id=_lowerCAmelCase    ,				**_lowerCAmelCase    ,				)






class     __lowerCAmelCase     ( a					):






  """simple docstring"""





  @property
  def 					lowerCAmelCase__  (     self       :       Tuple  )   -> Mapping[str, Mapping[int, str]]:




     """simple docstring"""
     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       :       Dict  )   -> Mapping[str, Mapping[int, str]]:




     """simple docstring"""
     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       :       str    ,				_lowerCAmelCase       :       PreTrainedTokenizer    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       bool = False    ,				_lowerCAmelCase       :       Optional[TensorType] = None    ,				)   -> Mapping[str, Any]:




     """simple docstring"""
     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       :       Optional[int]    ,				_lowerCAmelCase       :       PreTrainedTokenizer    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       bool = False    ,				_lowerCAmelCase       :       Optional[TensorType] = None    ,				)   -> Mapping[str, Any]:




     """simple docstring"""
     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       :       Optional[Any]    ,				_lowerCAmelCase       :       PreTrainedTokenizer    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       bool = False    ,				_lowerCAmelCase       :       Optional[TensorType] = None    ,				)   -> Mapping[str, Any]:




     """simple docstring"""
     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       :       Union[str, Any]    ,				_lowerCAmelCase       :       PreTrainedTokenizer    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       int = -1    ,				_lowerCAmelCase       :       bool = False    ,				_lowerCAmelCase       :       Optional[TensorType] = None    ,				)   -> Mapping[str, Any]:




     """simple docstring"""
     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       :       str    ,				_lowerCAmelCase       :       Union[str, Any]    ,				_lowerCAmelCase       :       Dict    ,				_lowerCAmelCase       :       Union[str, Any]    ,				_lowerCAmelCase       :       Dict  )   -> str:




     """simple docstring"""
     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 typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


SCREAMING_SNAKE_CASE	:int  		=     logging.get_logger(__name__)

SCREAMING_SNAKE_CASE	:Optional[Any]  		=     {
    '''google/long-t5-local-base''': '''https://huggingface.co/google/long-t5-local-base/blob/main/config.json''',
    '''google/long-t5-local-large''': '''https://huggingface.co/google/long-t5-local-large/blob/main/config.json''',
    '''google/long-t5-tglobal-base''': '''https://huggingface.co/google/long-t5-tglobal-base/blob/main/config.json''',
    '''google/long-t5-tglobal-large''': '''https://huggingface.co/google/long-t5-tglobal-large/blob/main/config.json''',
}


class     __lowerCAmelCase     ( a					):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		    =       'longt5'
  _SCREAMING_SNAKE_CASE		    =       ['past_key_values']
  _SCREAMING_SNAKE_CASE		    =       {'hidden_size': 'd_model', 'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}



  def __init__(     self       :       int    ,				_lowerCAmelCase       :       str=3_2_1_2_8    ,				_lowerCAmelCase       :       Optional[Any]=5_1_2    ,				_lowerCAmelCase       :       List[str]=6_4    ,				_lowerCAmelCase       :       Optional[int]=2_0_4_8    ,				_lowerCAmelCase       :       Any=6    ,				_lowerCAmelCase       :       Any=None    ,				_lowerCAmelCase       :       str=8    ,				_lowerCAmelCase       :       List[str]=1_2_7    ,				_lowerCAmelCase       :       int=1_6    ,				_lowerCAmelCase       :       List[str]=3_2    ,				_lowerCAmelCase       :       Any=1_2_8    ,				_lowerCAmelCase       :       Union[str, Any]=0.1    ,				_lowerCAmelCase       :       Union[str, Any]=1e-6    ,				_lowerCAmelCase       :       str=1.0    ,				_lowerCAmelCase       :       Tuple="relu"    ,				_lowerCAmelCase       :       List[str]=True    ,				_lowerCAmelCase       :       Tuple="local"    ,				_lowerCAmelCase       :       Union[str, Any]=True    ,				_lowerCAmelCase       :       Optional[int]=0    ,				_lowerCAmelCase       :       str=1    ,				**_lowerCAmelCase       :       int    ,				)   -> Optional[int]:




     """simple docstring"""
     snake_case_					    =					vocab_size
     snake_case_					    =					d_model
     snake_case_					    =					d_kv
     snake_case_					    =					d_ff
     snake_case_					    =					num_layers
     # default = symmetry
     snake_case_					    =					num_decoder_layers if num_decoder_layers is not None else self.num_layers
     snake_case_					    =					num_heads
     snake_case_					    =					local_radius
     snake_case_					    =					global_block_size
     snake_case_					    =					relative_attention_num_buckets
     snake_case_					    =					relative_attention_max_distance
     snake_case_					    =					dropout_rate
     snake_case_					    =					layer_norm_epsilon
     snake_case_					    =					initializer_factor
     snake_case_					    =					feed_forward_proj
     snake_case_					    =					encoder_attention_type
     snake_case_					    =					use_cache

     snake_case_					    =					self.feed_forward_proj.split("-"  )
     snake_case_					    =					act_info[-1]
     snake_case_					    =					act_info[0] == "gated"

     if len(_lowerCAmelCase  ) > 1 and act_info[0] != "gated" or len(_lowerCAmelCase  ) > 2:
        raise ValueError(
            F'''`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.'''
            "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
            "'gated-gelu' or 'relu'"  )

     # for backwards compatibility
     if feed_forward_proj == "gated-gelu":
        snake_case_					    =					"gelu_new"

     super().__init__(
         pad_token_id=_lowerCAmelCase    ,				eos_token_id=_lowerCAmelCase    ,				is_encoder_decoder=_lowerCAmelCase    ,				**_lowerCAmelCase    ,				)






class     __lowerCAmelCase     ( a					):






  """simple docstring"""





  @property
  def 					lowerCAmelCase__  (     self       :       str  )   -> Mapping[str, Mapping[int, str]]:




     """simple docstring"""
     snake_case_					    =					{
         "input_ids": {0: "batch", 1: "encoder_sequence"},
         "attention_mask": {0: "batch", 1: "encoder_sequence"},
     }
     if self.use_past:
        snake_case_					    =					"past_encoder_sequence + sequence"
        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"  )

     return common_inputs



  @property
  def 					lowerCAmelCase__  (     self       :       List[Any]  )   -> int:




     """simple docstring"""
     return 1_3




