




'''simple docstring'''




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


a_   :	int							   = logging.get_logger(__name__)

a_   :	Union[str, Any]							   = {
    """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 						__UpperCamelCase (				lowerCamelCase__       ):
       lowercase			:					Dict   		='bigbird_pegasus'
       lowercase			:					Dict   		=['past_key_values']
       lowercase			:					List[str]   		={
           'num_attention_heads': 'encoder_attention_heads',
           'hidden_size': 'd_model',
           'attention_probs_dropout_prob': 'attention_dropout',
       }






       def __init__(							self,	lowerCAmelCase=96_103,	lowerCAmelCase=4_096,	lowerCAmelCase=16,	lowerCAmelCase=4_096,	lowerCAmelCase=16,	lowerCAmelCase=16,	lowerCAmelCase=4_096,	lowerCAmelCase=16,	lowerCAmelCase=0.0,	lowerCAmelCase=0.0,	lowerCAmelCase=True,	lowerCAmelCase=True,	lowerCAmelCase="gelu_new",	lowerCAmelCase=1_024,	lowerCAmelCase=0.1,	lowerCAmelCase=0.0,	lowerCAmelCase=0.0,	lowerCAmelCase=0.0_2,	lowerCAmelCase=2,	lowerCAmelCase=0.0,	lowerCAmelCase=True,	lowerCAmelCase=0,	lowerCAmelCase=2,	lowerCAmelCase=1,	lowerCAmelCase="block_sparse",	lowerCAmelCase=64,	lowerCAmelCase=3,	lowerCAmelCase=False,	**lowerCAmelCase,	):



              """simple docstring"""







              lowerCamelCase_							=vocab_size
              lowerCamelCase_							=max_position_embeddings
              lowerCamelCase_							=d_model
              lowerCamelCase_							=encoder_ffn_dim
              lowerCamelCase_							=encoder_layers
              lowerCamelCase_							=encoder_attention_heads
              lowerCamelCase_							=decoder_ffn_dim
              lowerCamelCase_							=decoder_layers
              lowerCamelCase_							=decoder_attention_heads
              lowerCamelCase_							=dropout
              lowerCamelCase_							=attention_dropout
              lowerCamelCase_							=activation_dropout
              lowerCamelCase_							=activation_function
              lowerCamelCase_							=init_std
              lowerCamelCase_							=encoder_layerdrop
              lowerCamelCase_							=decoder_layerdrop
              lowerCamelCase_							=classifier_dropout
              lowerCamelCase_							=use_cache
              lowerCamelCase_							=encoder_layers
              lowerCamelCase_							=scale_embedding  # scale factor will be sqrt(d_model) if True

              # extra config
              lowerCamelCase_							=attention_type
              lowerCamelCase_							=block_size
              lowerCamelCase_							=num_random_blocks
              lowerCamelCase_							=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 						__UpperCamelCase (				lowerCamelCase__       ):






       @property
       def   lowercase__  (							self						):



              """simple docstring"""







              if self.task in ["default", "seq2seq-lm"]:
                     lowerCamelCase_							=OrderedDict(
                         [
                             ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                             ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                         ]						)

                     if self.use_past:
                            lowerCamelCase_							={0: '''batch'''}
                            lowerCamelCase_							={0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
                     else:
                            lowerCamelCase_							={0: '''batch''', 1: '''decoder_sequence'''}
                            lowerCamelCase_							={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.
                     lowerCamelCase_							=OrderedDict(
                         [
                             ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                             ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                         ]						)
                     if self.use_past:
                            lowerCamelCase_,   lowerCamelCase_							=self.num_layers
                            for i in range(lowerCAmelCase						):
                                   lowerCamelCase_							={0: '''batch''', 2: '''past_sequence + sequence'''}
                                   lowerCamelCase_							={0: '''batch''', 2: '''past_sequence + sequence'''}
              else:
                     lowerCamelCase_							=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   lowercase__  (							self						):



              """simple docstring"""







              if self.task in ["default", "seq2seq-lm"]:
                     lowerCamelCase_							=super().outputs
              else:
                     lowerCamelCase_							=super(lowerCAmelCase,	self						).outputs
                     if self.use_past:
                            lowerCamelCase_,   lowerCamelCase_							=self.num_layers
                            for i in range(lowerCAmelCase						):
                                   lowerCamelCase_							={0: '''batch''', 2: '''past_sequence + sequence'''}
                                   lowerCamelCase_							={0: '''batch''', 2: '''past_sequence + sequence'''}
              return common_outputs






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = -1,	lowerCAmelCase = -1,	lowerCAmelCase = False,	lowerCAmelCase = None,	):



              """simple docstring"""







              lowerCamelCase_							=self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                  lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)

              # Generate decoder inputs
              lowerCamelCase_							=seq_length if not self.use_past else 1
              lowerCamelCase_							=self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                  lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)
              lowerCamelCase_							={f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
              lowerCamelCase_							=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
                     lowerCamelCase_,   lowerCamelCase_							=common_inputs['''input_ids'''].shape
                     lowerCamelCase_							=common_inputs['''decoder_input_ids'''].shape[1]
                     lowerCamelCase_,   lowerCamelCase_							=self.num_attention_heads
                     lowerCamelCase_							=(
                         batch,
                         num_encoder_attention_heads,
                         encoder_seq_length,
                         self._config.hidden_size // num_encoder_attention_heads,
                     )
                     lowerCamelCase_							=decoder_seq_length + 3
                     lowerCamelCase_							=(
                         batch,
                         num_decoder_attention_heads,
                         decoder_past_length,
                         self._config.hidden_size // num_decoder_attention_heads,
                     )

                     lowerCamelCase_							=torch.cat(
                         [common_inputs['''decoder_attention_mask'''], torch.ones(lowerCAmelCase,	lowerCAmelCase						)],	dim=1						)

                     lowerCamelCase_							=[]
                     # If the number of encoder and decoder layers are present in the model configuration, both are considered
                     lowerCamelCase_,   lowerCamelCase_							=self.num_layers
                     lowerCamelCase_							=min(lowerCAmelCase,	lowerCAmelCase						)
                     lowerCamelCase_							=max(lowerCAmelCase,	lowerCAmelCase						) - min_num_layers
                     lowerCamelCase_							='''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.
                     lowerCamelCase_							=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   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = -1,	lowerCAmelCase = -1,	lowerCAmelCase = False,	lowerCAmelCase = None,	):



              """simple docstring"""







              lowerCamelCase_							=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
                     lowerCamelCase_,   lowerCamelCase_							=common_inputs['''input_ids'''].shape
                     # Not using the same length for past_key_values
                     lowerCamelCase_							=seqlen + 2
                     lowerCamelCase_,   lowerCamelCase_							=self.num_layers
                     lowerCamelCase_,   lowerCamelCase_							=self.num_attention_heads
                     lowerCamelCase_							=(
                         batch,
                         num_encoder_attention_heads,
                         past_key_values_length,
                         self._config.hidden_size // num_encoder_attention_heads,
                     )

                     lowerCamelCase_							=common_inputs['''attention_mask'''].dtype
                     lowerCamelCase_							=torch.cat(
                         [common_inputs['''attention_mask'''], torch.ones(lowerCAmelCase,	lowerCAmelCase,	dtype=lowerCAmelCase						)],	dim=1						)
                     lowerCamelCase_							=[
                         (torch.zeros(lowerCAmelCase						), torch.zeros(lowerCAmelCase						)) for _ in range(lowerCAmelCase						)
                     ]
              return common_inputs






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = -1,	lowerCAmelCase = -1,	lowerCAmelCase = False,	lowerCAmelCase = None,	):



              """simple docstring"""







              lowerCamelCase_							=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
              lowerCamelCase_							=tokenizer.num_special_tokens_to_add(lowerCAmelCase						)
              lowerCamelCase_							=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
              lowerCamelCase_							=[''' '''.join([tokenizer.unk_token]						) * seq_length] * batch_size
              lowerCamelCase_							=dict(tokenizer(lowerCAmelCase,	return_tensors=lowerCAmelCase						)						)
              return common_inputs






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase = -1,	lowerCAmelCase = -1,	lowerCAmelCase = False,	lowerCAmelCase = None,	):



              """simple docstring"""







              if self.task in ["default", "seq2seq-lm"]:
                     lowerCamelCase_							=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":
                     lowerCamelCase_							=self._generate_dummy_inputs_for_causal_lm(
                         lowerCAmelCase,	batch_size=lowerCAmelCase,	seq_length=lowerCAmelCase,	is_pair=lowerCAmelCase,	framework=lowerCAmelCase						)
              else:
                     lowerCamelCase_							=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   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						):



              """simple docstring"""







              if self.task in ["default", "seq2seq-lm"]:
                     lowerCamelCase_							=super()._flatten_past_key_values_(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)
              else:
                     lowerCamelCase_							=super(lowerCAmelCase,	self						)._flatten_past_key_values_(
                         lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)










'''simple docstring'''




import warnings

from .generation import FlaxGenerationMixin


class 						__UpperCamelCase (				lowerCamelCase__       ):
       # warning at import time
       warnings.warn(
           'Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will '
           'be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead.'    ,      lowerCamelCase__    ,      )




