code: 
  
  
  
  
  
  
  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_				)
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  
  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__)
  
  
  
  
  
  