






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


__snake_case							      =						logging.get_logger(__name__)

__snake_case							      =						{
    """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_		(      lowercase				):



      """simple docstring"""
      UpperCamelCase_       :      str     ='bigbird_pegasus'
      UpperCamelCase_       :      List[str]     =['past_key_values']
      UpperCamelCase_       :      Optional[int]     ={
          'num_attention_heads': 'encoder_attention_heads',
          'hidden_size': 'd_model',
          'attention_probs_dropout_prob': 'attention_dropout',
      }
      def __init__(       self					,       SCREAMING_SNAKE_CASE_=9_6103					,       SCREAMING_SNAKE_CASE_=4096					,       SCREAMING_SNAKE_CASE_=16					,       SCREAMING_SNAKE_CASE_=4096					,       SCREAMING_SNAKE_CASE_=16					,       SCREAMING_SNAKE_CASE_=16					,       SCREAMING_SNAKE_CASE_=4096					,       SCREAMING_SNAKE_CASE_=16					,       SCREAMING_SNAKE_CASE_=0.0					,       SCREAMING_SNAKE_CASE_=0.0					,       SCREAMING_SNAKE_CASE_=True					,       SCREAMING_SNAKE_CASE_=True					,       SCREAMING_SNAKE_CASE_="gelu_new"					,       SCREAMING_SNAKE_CASE_=1024					,       SCREAMING_SNAKE_CASE_=0.1					,       SCREAMING_SNAKE_CASE_=0.0					,       SCREAMING_SNAKE_CASE_=0.0					,       SCREAMING_SNAKE_CASE_=0.02					,       SCREAMING_SNAKE_CASE_=2					,       SCREAMING_SNAKE_CASE_=0.0					,       SCREAMING_SNAKE_CASE_=True					,       SCREAMING_SNAKE_CASE_=0					,       SCREAMING_SNAKE_CASE_=2					,       SCREAMING_SNAKE_CASE_=1					,       SCREAMING_SNAKE_CASE_="block_sparse"					,       SCREAMING_SNAKE_CASE_=64					,       SCREAMING_SNAKE_CASE_=3					,       SCREAMING_SNAKE_CASE_=False					,       **SCREAMING_SNAKE_CASE_					,       )       ->   Any:
          UpperCamelCase :Dict										=     vocab_size
          UpperCamelCase :Optional[Any]										=     max_position_embeddings
          UpperCamelCase :List[str]										=     d_model
          UpperCamelCase :str										=     encoder_ffn_dim
          UpperCamelCase :List[Any]										=     encoder_layers
          UpperCamelCase :Any										=     encoder_attention_heads
          UpperCamelCase :List[str]										=     decoder_ffn_dim
          UpperCamelCase :List[str]										=     decoder_layers
          UpperCamelCase :Any										=     decoder_attention_heads
          UpperCamelCase :Optional[Any]										=     dropout
          UpperCamelCase :Optional[int]										=     attention_dropout
          UpperCamelCase :Any										=     activation_dropout
          UpperCamelCase :Tuple										=     activation_function
          UpperCamelCase :Any										=     init_std
          UpperCamelCase :List[str]										=     encoder_layerdrop
          UpperCamelCase :Dict										=     decoder_layerdrop
          UpperCamelCase :Tuple										=     classifier_dropout
          UpperCamelCase :Optional[int]										=     use_cache
          UpperCamelCase :Optional[int]										=     encoder_layers
          UpperCamelCase :Union[str, Any]										=     scale_embedding  # scale factor will be sqrt(d_model) if True

          # extra config
          UpperCamelCase :int										=     attention_type
          UpperCamelCase :Any										=     block_size
          UpperCamelCase :Tuple										=     num_random_blocks
          UpperCamelCase :List[Any]										=     use_bias

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







class     UpperCAmelCase_		(      lowercase				):



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

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

              if self.use_past:
                  self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE_					,       direction='''inputs'''				)
          elif self.task == "causal-lm":
              # TODO: figure this case out.
              UpperCamelCase :List[Any]										=     OrderedDict(
                  [
                      ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
                      ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
                  ]				)
              if self.use_past:
                  UpperCamelCase							,     UpperCamelCase :Any										=     self.num_layers
                  for i in range(SCREAMING_SNAKE_CASE_				):
                      UpperCamelCase :Optional[Any]										=     {0: '''batch''', 2: '''past_sequence + sequence'''}
                      UpperCamelCase :int										=     {0: '''batch''', 2: '''past_sequence + sequence'''}
          else:
              UpperCamelCase :Optional[Any]										=     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]]:
          if self.task in ["default", "seq2seq-lm"]:
              UpperCamelCase :Tuple										=     super().outputs
          else:
              UpperCamelCase :Tuple										=     super(SCREAMING_SNAKE_CASE_					,       self				).outputs
              if self.use_past:
                  UpperCamelCase							,     UpperCamelCase :Union[str, Any]										=     self.num_layers
                  for i in range(SCREAMING_SNAKE_CASE_				):
                      UpperCamelCase :Optional[Any]										=     {0: '''batch''', 2: '''past_sequence + sequence'''}
                      UpperCamelCase :str										=     {0: '''batch''', 2: '''past_sequence + sequence'''}
          return common_outputs
      def 					UpperCAmelCase		(       self					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = False					,       SCREAMING_SNAKE_CASE_ = None					,       )       ->   Mapping[str, Any]:
          UpperCamelCase :Tuple										=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_				)

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

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

              UpperCamelCase :Tuple										=     torch.cat(
                  [common_inputs['''decoder_attention_mask'''], torch.ones(SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_				)]					,       dim=1				)

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

              for _ in range(SCREAMING_SNAKE_CASE_				):
                  common_inputs["past_key_values"].append(
                      (
                          torch.zeros(SCREAMING_SNAKE_CASE_				),
                          torch.zeros(SCREAMING_SNAKE_CASE_				),
                          torch.zeros(SCREAMING_SNAKE_CASE_				),
                          torch.zeros(SCREAMING_SNAKE_CASE_				),
                      )				)
              # TODO: test this.
              UpperCamelCase :str										=     encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
              for _ in range(SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_				):
                  common_inputs["past_key_values"].append((torch.zeros(SCREAMING_SNAKE_CASE_				), torch.zeros(SCREAMING_SNAKE_CASE_				))				)
          return common_inputs
      def 					UpperCAmelCase		(       self					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = False					,       SCREAMING_SNAKE_CASE_ = None					,       )       ->   Mapping[str, Any]:
          UpperCamelCase :Optional[Any]										=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_				)

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

              UpperCamelCase :Optional[int]										=     common_inputs['''attention_mask'''].dtype
              UpperCamelCase :Union[str, Any]										=     torch.cat(
                  [common_inputs['''attention_mask'''], torch.ones(SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       dtype=SCREAMING_SNAKE_CASE_				)]					,       dim=1				)
              UpperCamelCase :Dict										=     [
                  (torch.zeros(SCREAMING_SNAKE_CASE_				), torch.zeros(SCREAMING_SNAKE_CASE_				)) for _ in range(SCREAMING_SNAKE_CASE_				)
              ]
          return common_inputs
      def 					UpperCAmelCase		(       self					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = False					,       SCREAMING_SNAKE_CASE_ = 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
          UpperCamelCase :str										=     compute_effective_axis_dimension(
              SCREAMING_SNAKE_CASE_					,       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
          UpperCamelCase :List[str]										=     tokenizer.num_special_tokens_to_add(SCREAMING_SNAKE_CASE_				)
          UpperCamelCase :List[Any]										=     compute_effective_axis_dimension(
              SCREAMING_SNAKE_CASE_					,       fixed_dimension=OnnxConfig.default_fixed_sequence					,       num_token_to_add=SCREAMING_SNAKE_CASE_				)

          # Generate dummy inputs according to compute batch and sequence
          UpperCamelCase :Optional[int]										=     [''' '''.join([tokenizer.unk_token]				) * seq_length] * batch_size
          UpperCamelCase :List[Any]										=     dict(tokenizer(SCREAMING_SNAKE_CASE_					,       return_tensors=SCREAMING_SNAKE_CASE_				)				)
          return common_inputs
      def 					UpperCAmelCase		(       self					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = -1					,       SCREAMING_SNAKE_CASE_ = False					,       SCREAMING_SNAKE_CASE_ = None					,       )       ->   Mapping[str, Any]:
          if self.task in ["default", "seq2seq-lm"]:
              UpperCamelCase :Any										=     self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                  SCREAMING_SNAKE_CASE_					,       batch_size=SCREAMING_SNAKE_CASE_					,       seq_length=SCREAMING_SNAKE_CASE_					,       is_pair=SCREAMING_SNAKE_CASE_					,       framework=SCREAMING_SNAKE_CASE_				)

          elif self.task == "causal-lm":
              UpperCamelCase :int										=     self._generate_dummy_inputs_for_causal_lm(
                  SCREAMING_SNAKE_CASE_					,       batch_size=SCREAMING_SNAKE_CASE_					,       seq_length=SCREAMING_SNAKE_CASE_					,       is_pair=SCREAMING_SNAKE_CASE_					,       framework=SCREAMING_SNAKE_CASE_				)
          else:
              UpperCamelCase :Tuple										=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                  SCREAMING_SNAKE_CASE_					,       batch_size=SCREAMING_SNAKE_CASE_					,       seq_length=SCREAMING_SNAKE_CASE_					,       is_pair=SCREAMING_SNAKE_CASE_					,       framework=SCREAMING_SNAKE_CASE_				)

          return common_inputs




      def 					UpperCAmelCase		(       self					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_				)       ->   List[str]:
          if self.task in ["default", "seq2seq-lm"]:
              UpperCamelCase :Any										=     super()._flatten_past_key_values_(SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_				)
          else:
              UpperCamelCase :List[str]										=     super(SCREAMING_SNAKE_CASE_					,       self				)._flatten_past_key_values_(
                  SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_					,       SCREAMING_SNAKE_CASE_				)













from typing import TYPE_CHECKING

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


__snake_case							      =						{
    """configuration_open_llama""": ["""OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """OpenLlamaConfig"""],
}

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

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

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case							      =						[
     """OpenLlamaForCausalLM""",
     """OpenLlamaModel""",
     """OpenLlamaPreTrainedModel""",
     """OpenLlamaForSequenceClassification""",
 ]


if TYPE_CHECKING:
 from .configuration_open_llama import OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenLlamaConfig

 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from transformers import LlamaTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from transformers import LlamaTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_open_llama import (
      OpenLlamaForCausalLM,
      OpenLlamaForSequenceClassification,
      OpenLlamaModel,
      OpenLlamaPreTrainedModel,
  )


else:
 import sys

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





