






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


UpperCamelCase											=     logging.get_logger(__name__)

UpperCamelCase											=     {
    """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  lowercase_     (_UpperCAmelCase  ):
      A__						:		str							      =					'''bigbird_pegasus'''
      A__						:		str							      =					['''past_key_values''']
      A__						:		Any							      =					{
          '''num_attention_heads''': '''encoder_attention_heads''',
          '''hidden_size''': '''d_model''',
          '''attention_probs_dropout_prob''': '''attention_dropout''',
      }






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







        '''simple docstring'''


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

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

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







class  lowercase_     (_UpperCAmelCase  ):






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







        '''simple docstring'''


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

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

          if self.use_past:
            self.fill_with_past_key_values_(a_  ,  direction="inputs"						)
        elif self.task == "causal-lm":
          # TODO: figure this case out.
          _a					      =						OrderedDict(
              [
                  ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                  ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
              ]						)
          if self.use_past:
            _a      ,						_a					      =						self.num_layers
            for i in range(a_						):
              _a					      =						{0: "batch", 2: "past_sequence + sequence"}
              _a					      =						{0: "batch", 2: "past_sequence + sequence"}
        else:
          _a					      =						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]]:







        '''simple docstring'''


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






      def       lowerCamelCase__     (     self  ,  a_  ,  a_ = -1  ,  a_ = -1  ,  a_ = False  ,  a_ = None  ,  )						->Mapping[str, Any]:







        '''simple docstring'''


        _a					      =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            a_  ,  a_  ,  a_  ,  a_  ,  a_						)

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

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

          _a					      =						torch.cat(
              [common_inputs["decoder_attention_mask"], torch.ones(a_  ,  a_						)]  ,  dim=1						)

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

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






      def       lowerCamelCase__     (     self  ,  a_  ,  a_ = -1  ,  a_ = -1  ,  a_ = False  ,  a_ = None  ,  )						->Mapping[str, Any]:







        '''simple docstring'''


        _a					      =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            a_  ,  a_  ,  a_  ,  a_  ,  a_						)

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

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






      def       lowerCamelCase__     (     self  ,  a_  ,  a_ = -1  ,  a_ = -1  ,  a_ = False  ,  a_ = None  ,  )						->Mapping[str, Any]:







        '''simple docstring'''


        _a					      =						compute_effective_axis_dimension(
            a_  ,  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
        _a					      =						tokenizer.num_special_tokens_to_add(a_						)
        _a					      =						compute_effective_axis_dimension(
            a_  ,  fixed_dimension=OnnxConfig.default_fixed_sequence  ,  num_token_to_add=a_						)

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






      def       lowerCamelCase__     (     self  ,  a_  ,  a_ = -1  ,  a_ = -1  ,  a_ = False  ,  a_ = None  ,  )						->Mapping[str, Any]:







        '''simple docstring'''


        if self.task in ["default", "seq2seq-lm"]:
          _a					      =						self._generate_dummy_inputs_for_default_and_seqaseq_lm(
              a_  ,  batch_size=a_  ,  seq_length=a_  ,  is_pair=a_  ,  framework=a_						)

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

        return common_inputs






      def       lowerCamelCase__     (     self  ,  a_  ,  a_  ,  a_  ,  a_						)						->str:







        '''simple docstring'''


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










"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase											=     {"""configuration_pegasus""": ["""PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP""", """PegasusConfig"""]}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     ["""PegasusTokenizer"""]

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     ["""PegasusTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """PegasusForCausalLM""",
         """PegasusForConditionalGeneration""",
         """PegasusModel""",
         """PegasusPreTrainedModel""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """TFPegasusForConditionalGeneration""",
         """TFPegasusModel""",
         """TFPegasusPreTrainedModel""",
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase											=     [
         """FlaxPegasusForConditionalGeneration""",
         """FlaxPegasusModel""",
         """FlaxPegasusPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_pegasus import PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusConfig

     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_pegasus import PegasusTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_pegasus_fast import PegasusTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_pegasus import (
              PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
              PegasusForCausalLM,
              PegasusForConditionalGeneration,
              PegasusModel,
              PegasusPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel, TFPegasusPreTrainedModel

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_pegasus import (
              FlaxPegasusForConditionalGeneration,
              FlaxPegasusModel,
              FlaxPegasusPreTrainedModel,
          )

else:
     import sys

     UpperCamelCase											=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


