

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


lowerCamelCase     =      logging.get_logger(__name__)

lowerCamelCase     =      {
    '''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 		__magic_name__    (							lowerCamelCase__							):

   '''simple docstring'''




   lowerCamelCase__				:			Optional[int]												=     'bigbird_pegasus'
   lowerCamelCase__				:			str												=     ['past_key_values']
   lowerCamelCase__				:			Dict												=     {
       'num_attention_heads': 'encoder_attention_heads',
       'hidden_size': 'd_model',
       'attention_probs_dropout_prob': 'attention_dropout',
   }



   def __init__(       self,					lowercase_=96103,					lowercase_=4096,					lowercase_=16,					lowercase_=4096,					lowercase_=16,					lowercase_=16,					lowercase_=4096,					lowercase_=16,					lowercase_=0.0,					lowercase_=0.0,					lowercase_=True,					lowercase_=True,					lowercase_="gelu_new",					lowercase_=1024,					lowercase_=0.1,					lowercase_=0.0,					lowercase_=0.0,					lowercase_=0.02,					lowercase_=2,					lowercase_=0.0,					lowercase_=True,					lowercase_=0,					lowercase_=2,					lowercase_=1,					lowercase_="block_sparse",					lowercase_=64,					lowercase_=3,					lowercase_=False,					**lowercase_,					)							->					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=lowercase_,					bos_token_id=lowercase_,					eos_token_id=lowercase_,					is_encoder_decoder=lowercase_,					decoder_start_token_id=lowercase_,					**lowercase_,					)







class 		__magic_name__    (							lowerCamelCase__							):

   '''simple docstring'''




   @property
   def 				_UpperCAmelCase				(       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_(lowercase_,					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(lowercase_ ):
           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 				_UpperCAmelCase				(       self )							->					Mapping[str, Mapping[int, str]]:






     """simple docstring"""






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



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






     """simple docstring"""






     a__						=self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )

     # Generate decoder inputs
     a__						=seq_length if not self.use_past else 1
     a__						=self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )
     a__						={F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
     a__						=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
       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(lowercase_,					lowercase_ )],					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(lowercase_,					lowercase_ )
       a__						=max(lowercase_,					lowercase_ ) - min_num_layers
       a__						='''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.
       a__						=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,					lowercase_,					lowercase_ = -1,					lowercase_ = -1,					lowercase_ = False,					lowercase_ = None,					)							->					Mapping[str, Any]:






     """simple docstring"""






     a__						=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
       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(lowercase_,					lowercase_,					dtype=lowercase_ )],					dim=1 )
       a__						=[
           (torch.zeros(lowercase_ ), torch.zeros(lowercase_ )) for _ in range(lowercase_ )
       ]
     return common_inputs



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






     """simple docstring"""






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



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






     """simple docstring"""






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






     """simple docstring"""






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



import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           ErnieMForInformationExtraction,
           ErnieMForMultipleChoice,
           ErnieMForQuestionAnswering,
           ErnieMForSequenceClassification,
           ErnieMForTokenClassification,
           ErnieMModel,
       )
       from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST







class 		__magic_name__    :

   '''simple docstring'''




   def __init__(       self,					lowercase_,					lowercase_=13,					lowercase_=7,					lowercase_=True,					lowercase_=True,					lowercase_=True,					lowercase_=99,					lowercase_=32,					lowercase_=5,					lowercase_=4,					lowercase_=37,					lowercase_="gelu",					lowercase_=0.1,					lowercase_=0.1,					lowercase_=512,					lowercase_=16,					lowercase_=2,					lowercase_=0.02,					lowercase_=3,					lowercase_=4,					lowercase_=None,					)							->					List[Any]:






     """simple docstring"""






     a__						=parent
     a__						=batch_size
     a__						=seq_length
     a__						=is_training
     a__						=use_input_mask
     a__						=use_labels
     a__						=vocab_size
     a__						=hidden_size
     a__						=num_hidden_layers
     a__						=num_attention_heads
     a__						=intermediate_size
     a__						=hidden_act
     a__						=hidden_dropout_prob
     a__						=attention_probs_dropout_prob
     a__						=max_position_embeddings
     a__						=type_vocab_size
     a__						=type_sequence_label_size
     a__						=initializer_range
     a__						=num_labels
     a__						=num_choices
     a__						=scope



   def 				_UpperCAmelCase				(       self )							->					List[Any]:






     """simple docstring"""






     a__						=ids_tensor([self.batch_size, self.seq_length],					self.vocab_size )

     a__						=None
     if self.use_input_mask:
       a__						=random_attention_mask([self.batch_size, self.seq_length] )

     a__						=None
     a__						=None
     a__						=None
     if self.use_labels:
       a__						=ids_tensor([self.batch_size],					self.type_sequence_label_size )
       a__						=ids_tensor([self.batch_size, self.seq_length],					self.num_labels )
       a__						=ids_tensor([self.batch_size],					self.num_choices )

     a__						=self.get_config()

     return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



   def 				_UpperCAmelCase				(       self )							->					Optional[Any]:






     """simple docstring"""






     a__						=ids_tensor([self.batch_size, self.seq_length],					self.vocab_size )

     a__						=None
     if self.use_input_mask:
       a__						=random_attention_mask([self.batch_size, self.seq_length] )
     a__						=self.get_config()

     return config, input_ids, input_mask



   def 				_UpperCAmelCase				(       self )							->					Optional[int]:






     """simple docstring"""






     return ErnieMConfig(
         vocab_size=self.vocab_size,					hidden_size=self.hidden_size,					num_hidden_layers=self.num_hidden_layers,					num_attention_heads=self.num_attention_heads,					intermediate_size=self.intermediate_size,					hidden_act=self.hidden_act,					hidden_dropout_prob=self.hidden_dropout_prob,					attention_probs_dropout_prob=self.attention_probs_dropout_prob,					max_position_embeddings=self.max_position_embeddings,					type_vocab_size=self.type_vocab_size,					initializer_range=self.initializer_range,					)



   def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )							->					Union[str, Any]:






     """simple docstring"""






     a__						=ErnieMModel(config=lowercase_ )
     model.to(lowercase_ )
     model.eval()
     a__						=model(lowercase_,					return_dict=lowercase_ )
     self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size) )



   def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )							->					List[str]:






     """simple docstring"""






     a__						=ErnieMForQuestionAnswering(config=lowercase_ )
     model.to(lowercase_ )
     model.eval()
     a__						=model(
         lowercase_,					attention_mask=lowercase_,					start_positions=lowercase_,					end_positions=lowercase_,					)
     self.parent.assertEqual(result.start_logits.shape,					(self.batch_size, self.seq_length) )
     self.parent.assertEqual(result.end_logits.shape,					(self.batch_size, self.seq_length) )



   def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )							->					Dict:






     """simple docstring"""






     a__						=ErnieMForInformationExtraction(config=lowercase_ )
     model.to(lowercase_ )
     model.eval()
     a__						=torch.ones_like(lowercase_,					dtype=torch.floataa )
     a__						=model(
         lowercase_,					attention_mask=lowercase_,					start_positions=lowercase_,					end_positions=lowercase_,					)
     self.parent.assertEqual(result.start_logits.shape,					(self.batch_size, self.seq_length) )
     self.parent.assertEqual(result.end_logits.shape,					(self.batch_size, self.seq_length) )



   def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )							->					Union[str, Any]:






     """simple docstring"""






     a__						=self.num_labels
     a__						=ErnieMForSequenceClassification(lowercase_ )
     model.to(lowercase_ )
     model.eval()
     a__						=model(lowercase_,					attention_mask=lowercase_,					labels=lowercase_ )
     self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.num_labels) )



   def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )							->					Any:






     """simple docstring"""






     a__						=self.num_labels
     a__						=ErnieMForTokenClassification(config=lowercase_ )
     model.to(lowercase_ )
     model.eval()
     input_ids.to(lowercase_ )
     input_mask.to(lowercase_ )
     token_labels.to(lowercase_ )

     a__						=model(lowercase_,					attention_mask=lowercase_,					labels=lowercase_ )

     self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.seq_length, self.num_labels) )



   def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )							->					Optional[int]:






     """simple docstring"""






     a__						=self.num_choices
     a__						=ErnieMForMultipleChoice(config=lowercase_ )
     model.to(lowercase_ )
     model.eval()
     a__						=input_ids.unsqueeze(1 ).expand(-1,					self.num_choices,					-1 ).contiguous()
     a__						=input_mask.unsqueeze(1 ).expand(-1,					self.num_choices,					-1 ).contiguous()
     a__						=model(
         lowercase_,					attention_mask=lowercase_,					labels=lowercase_,					)
     self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.num_choices) )



   def 				_UpperCAmelCase				(       self )							->					Optional[int]:






     """simple docstring"""






     a__						=self.prepare_config_and_inputs()
     (
         (
         a__
     ),   (
         a__
     ),   (
         a__
     ),   (
         a__
     ),   (
         a__
     ),   (
         a__
     ),   
     )						=config_and_inputs
     a__						={'''input_ids''': input_ids, '''attention_mask''': input_mask}
     return config, inputs_dict







@require_torch
class 		__magic_name__    (							lowerCamelCase__	, lowerCamelCase__	, unittest.TestCase							):

   '''simple docstring'''




   lowerCamelCase__				:			List[Any]												=     (
       (
           ErnieMModel,
           ErnieMForMultipleChoice,
           ErnieMForQuestionAnswering,
           ErnieMForSequenceClassification,
           ErnieMForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowerCamelCase__				:			Union[str, Any]												=     ()
   lowerCamelCase__				:			Union[str, Any]												=     (
       {
           'feature-extraction': ErnieMModel,
           'question-answering': ErnieMForQuestionAnswering,
           'text-classification': ErnieMForSequenceClassification,
           'token-classification': ErnieMForTokenClassification,
           'zero-shot': ErnieMForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   lowerCamelCase__				:			Any												=     False



   def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ )							->					Optional[Any]:






     """simple docstring"""






     if pipeline_test_casse_name == "QAPipelineTests":
       return True

     return False



   def 				_UpperCAmelCase				(       self )							->					int:






     """simple docstring"""






     a__						=ErnieMModelTester(self )
     a__						=ConfigTester(self,					config_class=lowercase_,					hidden_size=37 )



   def 				_UpperCAmelCase				(       self )							->					Any:






     """simple docstring"""






     self.config_tester.run_common_tests()



   def 				_UpperCAmelCase				(       self )							->					Any:






     """simple docstring"""






     a__						=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_ )



   def 				_UpperCAmelCase				(       self )							->					Union[str, Any]:






     """simple docstring"""






     a__						=self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
       a__						=type
       self.model_tester.create_and_check_model(*lowercase_ )



   def 				_UpperCAmelCase				(       self )							->					Dict:






     """simple docstring"""






     a__						=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*lowercase_ )



   def 				_UpperCAmelCase				(       self )							->					Tuple:






     """simple docstring"""






     a__						=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*lowercase_ )



   def 				_UpperCAmelCase				(       self )							->					Any:






     """simple docstring"""






     a__						=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*lowercase_ )



   def 				_UpperCAmelCase				(       self )							->					str:






     """simple docstring"""






     a__						=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_information_extraction(*lowercase_ )



   def 				_UpperCAmelCase				(       self )							->					Any:






     """simple docstring"""






     a__						=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*lowercase_ )



   @slow
   def 				_UpperCAmelCase				(       self )							->					str:






     """simple docstring"""






     for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       a__						=ErnieMModel.from_pretrained(lowercase_ )
       self.assertIsNotNone(lowercase_ )







@require_torch
class 		__magic_name__    (							unittest.TestCase							):

   '''simple docstring'''




   @slow
   def 				_UpperCAmelCase				(       self )							->					str:






     """simple docstring"""






     a__						=ErnieMModel.from_pretrained('''susnato/ernie-m-base_pytorch''' )
     model.eval()
     a__						=torch.tensor([[0, 1, 2, 3, 4, 5]] )
     a__						=model(lowercase_ )[0]

     # TODO Replace vocab size
     a__						=768

     a__						=torch.Size((1, 6, hidden_size) )
     self.assertEqual(output.shape,					lowercase_ )

     a__						=torch.tensor(
         [[[-0.0012, 0.1245, -0.0214], [-0.0742, 0.0244, -0.0771], [-0.0333, 0.1164, -0.1554]]] )

     self.assertTrue(torch.allclose(output[:, :3, :3],					lowercase_,					atol=1E-3 ) )
