





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 			_lowerCamelCase	(      UpperCamelCase_      ):



      """simple docstring"""






      UpperCAmelCase_  : int				      ='bigbird_pegasus'
      UpperCAmelCase_  : Tuple				      =['past_key_values']
      UpperCAmelCase_  : str				      ={
          'num_attention_heads': 'encoder_attention_heads',
          'hidden_size': 'd_model',
          'attention_probs_dropout_prob': 'attention_dropout',
      }







      def __init__(	self			,  UpperCAmelCase=96103			,  UpperCAmelCase=4096			,  UpperCAmelCase=16			,  UpperCAmelCase=4096			,  UpperCAmelCase=16			,  UpperCAmelCase=16			,  UpperCAmelCase=4096			,  UpperCAmelCase=16			,  UpperCAmelCase=0.0			,  UpperCAmelCase=0.0			,  UpperCAmelCase=True			,  UpperCAmelCase=True			,  UpperCAmelCase="gelu_new"			,  UpperCAmelCase=1024			,  UpperCAmelCase=0.1			,  UpperCAmelCase=0.0			,  UpperCAmelCase=0.0			,  UpperCAmelCase=0.02			,  UpperCAmelCase=2			,  UpperCAmelCase=0.0			,  UpperCAmelCase=True			,  UpperCAmelCase=0			,  UpperCAmelCase=2			,  UpperCAmelCase=1			,  UpperCAmelCase="block_sparse"			,  UpperCAmelCase=64			,  UpperCAmelCase=3			,  UpperCAmelCase=False			,  **UpperCAmelCase			,  )    ->				str:






        '''simple docstring'''






        __snake_case				:			Optional[int]            =    vocab_size
        __snake_case				:			Tuple            =    max_position_embeddings
        __snake_case				:			Any            =    d_model
        __snake_case				:			Optional[Any]            =    encoder_ffn_dim
        __snake_case				:			Union[str, Any]            =    encoder_layers
        __snake_case				:			int            =    encoder_attention_heads
        __snake_case				:			Optional[Any]            =    decoder_ffn_dim
        __snake_case				:			int            =    decoder_layers
        __snake_case				:			str            =    decoder_attention_heads
        __snake_case				:			Union[str, Any]            =    dropout
        __snake_case				:			str            =    attention_dropout
        __snake_case				:			int            =    activation_dropout
        __snake_case				:			Union[str, Any]            =    activation_function
        __snake_case				:			Optional[Any]            =    init_std
        __snake_case				:			Optional[Any]            =    encoder_layerdrop
        __snake_case				:			int            =    decoder_layerdrop
        __snake_case				:			Optional[int]            =    classifier_dropout
        __snake_case				:			int            =    use_cache
        __snake_case				:			Dict            =    encoder_layers
        __snake_case				:			List[str]            =    scale_embedding  # scale factor will be sqrt(d_model) if True

        # extra config
        __snake_case				:			Tuple            =    attention_type
        __snake_case				:			str            =    block_size
        __snake_case				:			Optional[int]            =    num_random_blocks
        __snake_case				:			Optional[int]            =    use_bias

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




class 			_lowerCamelCase	(      UpperCamelCase_      ):



      """simple docstring"""






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






        '''simple docstring'''






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

          if self.use_past:
            __snake_case				:			str            =    {0: 'batch'}
            __snake_case				:			Optional[int]            =    {0: 'batch', 1: 'past_decoder_sequence + sequence'}
          else:
            __snake_case				:			Optional[Any]            =    {0: 'batch', 1: 'decoder_sequence'}
            __snake_case				:			List[Any]            =    {0: 'batch', 1: 'decoder_sequence'}

          if self.use_past:
            self.fill_with_past_key_values_(lowercase_			,  direction="inputs"				)
        elif self.task == "causal-lm":
          # TODO: figure this case out.
          __snake_case				:			Union[str, Any]            =    OrderedDict(
              [
                  ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                  ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
              ]				)
          if self.use_past:
            __snake_case				:			Union[str, Any]            =    self.num_layers
            for i in range(lowercase_				):
              __snake_case				:			Dict            =    {0: 'batch', 2: 'past_sequence + sequence'}
              __snake_case				:			List[str]            =    {0: 'batch', 2: 'past_sequence + sequence'}
        else:
          __snake_case				:			Optional[int]            =    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   UpperCAmelCase						(	self				)    ->				Mapping[str, Mapping[int, str]]:






        '''simple docstring'''






        if self.task in ["default", "seq2seq-lm"]:
          __snake_case				:			Dict            =    super().outputs
        else:
          __snake_case				:			Optional[Any]            =    super(lowercase_			,  self				).outputs
          if self.use_past:
            __snake_case				:			str            =    self.num_layers
            for i in range(lowercase_				):
              __snake_case				:			Tuple            =    {0: 'batch', 2: 'past_sequence + sequence'}
              __snake_case				:			List[Any]            =    {0: 'batch', 2: 'past_sequence + sequence'}
        return common_outputs







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






        '''simple docstring'''






        __snake_case				:			int            =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowercase_			,  lowercase_			,  lowercase_			,  lowercase_			,  lowercase_				)

        # Generate decoder inputs
        __snake_case				:			str            =    seq_length if not self.use_past else 1
        __snake_case				:			List[Any]            =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowercase_			,  lowercase_			,  lowercase_			,  lowercase_			,  lowercase_				)
        __snake_case				:			List[Any]            =    {F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
        __snake_case				:			Union[str, Any]            =    dict(**lowercase_			,  **lowercase_				)

        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				:			int            =    common_inputs['input_ids'].shape
          __snake_case				:			List[Any]            =    common_inputs['decoder_input_ids'].shape[1]
          __snake_case				:			Union[str, Any]            =    self.num_attention_heads
          __snake_case				:			Dict            =    (
              batch,
              num_encoder_attention_heads,
              encoder_seq_length,
              self._config.hidden_size // num_encoder_attention_heads,
          )
          __snake_case				:			List[Any]            =    decoder_seq_length + 3
          __snake_case				:			int            =    (
              batch,
              num_decoder_attention_heads,
              decoder_past_length,
              self._config.hidden_size // num_decoder_attention_heads,
          )

          __snake_case				:			Union[str, Any]            =    torch.cat(
              [common_inputs["decoder_attention_mask"], torch.ones(lowercase_			,  lowercase_				)]			,  dim=1				)

          __snake_case				:			Optional[Any]            =    []
          # If the number of encoder and decoder layers are present in the model configuration, both are considered
          __snake_case				:			Tuple            =    self.num_layers
          __snake_case				:			Optional[int]            =    min(lowercase_			,  lowercase_				)
          __snake_case				:			int            =    max(lowercase_			,  lowercase_				) - min_num_layers
          __snake_case				:			str            =    'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

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







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






        '''simple docstring'''






        __snake_case				:			str            =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowercase_			,  lowercase_			,  lowercase_			,  lowercase_			,  lowercase_				)

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

          __snake_case				:			Tuple            =    common_inputs['attention_mask'].dtype
          __snake_case				:			Tuple            =    torch.cat(
              [common_inputs["attention_mask"], torch.ones(lowercase_			,  lowercase_			,  dtype=lowercase_				)]			,  dim=1				)
          __snake_case				:			Dict            =    [
              (torch.zeros(lowercase_				), torch.zeros(lowercase_				)) for _ in range(lowercase_				)
          ]
        return common_inputs







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






        '''simple docstring'''






        __snake_case				:			Optional[Any]            =    compute_effective_axis_dimension(
            lowercase_			,  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				:			Tuple            =    tokenizer.num_special_tokens_to_add(lowercase_				)
        __snake_case				:			Any            =    compute_effective_axis_dimension(
            lowercase_			,  fixed_dimension=OnnxConfig.default_fixed_sequence			,  num_token_to_add=lowercase_				)

        # Generate dummy inputs according to compute batch and sequence
        __snake_case				:			Any            =    [' '.join([tokenizer.unk_token]				) * seq_length] * batch_size
        __snake_case				:			int            =    dict(tokenizer(lowercase_			,  return_tensors=lowercase_				)				)
        return common_inputs







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






        '''simple docstring'''






        if self.task in ["default", "seq2seq-lm"]:
          __snake_case				:			Union[str, Any]            =    self._generate_dummy_inputs_for_default_and_seqaseq_lm(
              lowercase_			,  batch_size=lowercase_			,  seq_length=lowercase_			,  is_pair=lowercase_			,  framework=lowercase_				)

        elif self.task == "causal-lm":
          __snake_case				:			List[str]            =    self._generate_dummy_inputs_for_causal_lm(
              lowercase_			,  batch_size=lowercase_			,  seq_length=lowercase_			,  is_pair=lowercase_			,  framework=lowercase_				)
        else:
          __snake_case				:			List[Any]            =    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              lowercase_			,  batch_size=lowercase_			,  seq_length=lowercase_			,  is_pair=lowercase_			,  framework=lowercase_				)

        return common_inputs







      def   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase			,  UpperCAmelCase				)    ->				Dict:






        '''simple docstring'''






        if self.task in ["default", "seq2seq-lm"]:
          __snake_case				:			Union[str, Any]            =    super()._flatten_past_key_values_(lowercase_			,  lowercase_			,  lowercase_			,  lowercase_				)
        else:
          __snake_case				:			List[str]            =    super(lowercase_			,  self				)._flatten_past_key_values_(
              lowercase_			,  lowercase_			,  lowercase_			,  lowercase_				)








from typing import TYPE_CHECKING

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


_UpperCamelCase    =	{
    '''configuration_dpr''': ['''DPR_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''DPRConfig'''],
    '''tokenization_dpr''': [
        '''DPRContextEncoderTokenizer''',
        '''DPRQuestionEncoderTokenizer''',
        '''DPRReaderOutput''',
        '''DPRReaderTokenizer''',
    ],
}


try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''DPRContextEncoderTokenizerFast''',
         '''DPRQuestionEncoderTokenizerFast''',
         '''DPRReaderTokenizerFast''',
     ]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''DPRContextEncoder''',
         '''DPRPretrainedContextEncoder''',
         '''DPRPreTrainedModel''',
         '''DPRPretrainedQuestionEncoder''',
         '''DPRPretrainedReader''',
         '''DPRQuestionEncoder''',
         '''DPRReader''',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCamelCase    =	[
         '''TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''TFDPRContextEncoder''',
         '''TFDPRPretrainedContextEncoder''',
         '''TFDPRPretrainedQuestionEncoder''',
         '''TFDPRPretrainedReader''',
         '''TFDPRQuestionEncoder''',
         '''TFDPRReader''',
     ]


if TYPE_CHECKING:
     from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
     from .tokenization_dpr import (
         DPRContextEncoderTokenizer,
         DPRQuestionEncoderTokenizer,
         DPRReaderOutput,
         DPRReaderTokenizer,
     )

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_dpr_fast import (
              DPRContextEncoderTokenizerFast,
              DPRQuestionEncoderTokenizerFast,
              DPRReaderTokenizerFast,
          )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_dpr import (
              DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPRContextEncoder,
              DPRPretrainedContextEncoder,
              DPRPreTrainedModel,
              DPRPretrainedQuestionEncoder,
              DPRPretrainedReader,
              DPRQuestionEncoder,
              DPRReader,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_dpr import (
              TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFDPRContextEncoder,
              TFDPRPretrainedContextEncoder,
              TFDPRPretrainedQuestionEncoder,
              TFDPRPretrainedReader,
              TFDPRQuestionEncoder,
              TFDPRReader,
          )

else:
     import sys

     _UpperCamelCase    =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

