






'''simple docstring'''







import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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 (
							    MODEL_FOR_PRETRAINING_MAPPING,
							    MegatronBertForCausalLM,
							    MegatronBertForMaskedLM,
							    MegatronBertForMultipleChoice,
							    MegatronBertForNextSentencePrediction,
							    MegatronBertForPreTraining,
							    MegatronBertForQuestionAnswering,
							    MegatronBertForSequenceClassification,
							    MegatronBertForTokenClassification,
							    MegatronBertModel,
							)




class 		UpperCamelCase__     :




		'''simple docstring'''





		def __init__(   self ,_lowerCAmelCase ,_lowerCAmelCase=13 ,_lowerCAmelCase=7 ,_lowerCAmelCase=True ,_lowerCAmelCase=True ,_lowerCAmelCase=True ,_lowerCAmelCase=True ,_lowerCAmelCase=99 ,_lowerCAmelCase=64 ,_lowerCAmelCase=32 ,_lowerCAmelCase=5 ,_lowerCAmelCase=4 ,_lowerCAmelCase=37 ,_lowerCAmelCase="gelu" ,_lowerCAmelCase=0.1 ,_lowerCAmelCase=0.1 ,_lowerCAmelCase=5_12 ,_lowerCAmelCase=16 ,_lowerCAmelCase=2 ,_lowerCAmelCase=0.02 ,_lowerCAmelCase=3 ,_lowerCAmelCase=4 ,_lowerCAmelCase=None ,):
								lowerCamelCase__								=	parent
								lowerCamelCase__								=	batch_size
								lowerCamelCase__								=	seq_length
								lowerCamelCase__								=	is_training
								lowerCamelCase__								=	use_input_mask
								lowerCamelCase__								=	use_token_type_ids
								lowerCamelCase__								=	use_labels
								lowerCamelCase__								=	vocab_size
								lowerCamelCase__								=	hidden_size
								lowerCamelCase__								=	embedding_size
								lowerCamelCase__								=	num_hidden_layers
								lowerCamelCase__								=	num_attention_heads
								lowerCamelCase__								=	intermediate_size
								lowerCamelCase__								=	hidden_act
								lowerCamelCase__								=	hidden_dropout_prob
								lowerCamelCase__								=	attention_probs_dropout_prob
								lowerCamelCase__								=	max_position_embeddings
								lowerCamelCase__								=	type_vocab_size
								lowerCamelCase__								=	type_sequence_label_size
								lowerCamelCase__								=	initializer_range
								lowerCamelCase__								=	num_labels
								lowerCamelCase__								=	num_choices
								lowerCamelCase__								=	scope






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size							)

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

								lowerCamelCase__								=	None
								if self.use_token_type_ids:
														lowerCamelCase__								=	ids_tensor([self.batch_size, self.seq_length] ,self.type_vocab_size							)

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

								lowerCamelCase__								=	self.get_config()

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






		def       UpperCamelCase_       (   self							):
								return MegatronBertConfig(
								    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 ,embedding_size=self.embedding_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 ,is_decoder=_lowerCAmelCase ,initializer_range=self.initializer_range ,)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	MegatronBertModel(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase							)
								lowerCamelCase__								=	model(_lowerCAmelCase ,token_type_ids=_lowerCAmelCase							)
								lowerCamelCase__								=	model(_lowerCAmelCase							)

								self.parent.assertEqual(result.last_hidden_state.shape ,(self.batch_size, self.seq_length, self.hidden_size)							)
								self.parent.assertEqual(result.pooler_output.shape ,(self.batch_size, self.hidden_size)							)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	MegatronBertForMaskedLM(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,labels=_lowerCAmelCase							)
								self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.vocab_size)							)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	MegatronBertForCausalLM(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,labels=_lowerCAmelCase							)
								self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.vocab_size)							)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	MegatronBertForNextSentencePrediction(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(
								    _lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,labels=_lowerCAmelCase ,)
								self.parent.assertEqual(result.logits.shape ,(self.batch_size, 2)							)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	MegatronBertForPreTraining(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(
								    _lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,labels=_lowerCAmelCase ,next_sentence_label=_lowerCAmelCase ,)
								self.parent.assertEqual(result.prediction_logits.shape ,(self.batch_size, self.seq_length, self.vocab_size)							)
								self.parent.assertEqual(result.seq_relationship_logits.shape ,(self.batch_size, 2)							)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	MegatronBertForQuestionAnswering(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(
								    _lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,start_positions=_lowerCAmelCase ,end_positions=_lowerCAmelCase ,)
								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 ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	self.num_labels
								lowerCamelCase__								=	MegatronBertForSequenceClassification(_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,labels=_lowerCAmelCase							)
								self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_labels)							)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	self.num_labels
								lowerCamelCase__								=	MegatronBertForTokenClassification(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,labels=_lowerCAmelCase							)
								self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.num_labels)							)






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	self.num_choices
								lowerCamelCase__								=	MegatronBertForMultipleChoice(config=_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.eval()
								lowerCamelCase__								=	input_ids.unsqueeze(1							).expand(-1 ,self.num_choices ,-1							).contiguous()
								lowerCamelCase__								=	token_type_ids.unsqueeze(1							).expand(-1 ,self.num_choices ,-1							).contiguous()
								lowerCamelCase__								=	input_mask.unsqueeze(1							).expand(-1 ,self.num_choices ,-1							).contiguous()
								lowerCamelCase__								=	model(
								    _lowerCAmelCase ,attention_mask=_lowerCAmelCase ,token_type_ids=_lowerCAmelCase ,labels=_lowerCAmelCase ,)
								self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_choices)							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase__      
								)							,      (
								    lowerCamelCase__      
								)							,      (
								    lowerCamelCase__      
								)							,      (
								    lowerCamelCase__      
								)							,      (
								    lowerCamelCase__      
								)							,      (
								    lowerCamelCase__      
								)							,      (
								    lowerCamelCase__      
								)							,            
								)								=	config_and_inputs
								lowerCamelCase__								=	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
								return config, inputs_dict




@require_torch
class 		UpperCamelCase__     (a ,a ,unittest.TestCase      ):




		'''simple docstring'''





		_UpperCamelCase						=       (
		    (
		        MegatronBertModel,
		        MegatronBertForMaskedLM,
		        MegatronBertForCausalLM,
		        MegatronBertForMultipleChoice,
		        MegatronBertForNextSentencePrediction,
		        MegatronBertForPreTraining,
		        MegatronBertForQuestionAnswering,
		        MegatronBertForSequenceClassification,
		        MegatronBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		_UpperCamelCase						=       (
		    {
		        "feature-extraction": MegatronBertModel,
		        "fill-mask": MegatronBertForMaskedLM,
		        "question-answering": MegatronBertForQuestionAnswering,
		        "text-classification": MegatronBertForSequenceClassification,
		        "text-generation": MegatronBertForCausalLM,
		        "token-classification": MegatronBertForTokenClassification,
		        "zero-shot": MegatronBertForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		_UpperCamelCase						=       True
		# test_resize_embeddings = False
		_UpperCamelCase						=       False






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase=False							):
								lowerCamelCase__								=	super()._prepare_for_class(_lowerCAmelCase ,_lowerCAmelCase ,return_labels=_lowerCAmelCase							)

								if return_labels:
														if model_class in get_values(_lowerCAmelCase							):
																				lowerCamelCase__								=	torch.zeros(
																				    (self.model_tester.batch_size, self.model_tester.seq_length) ,dtype=torch.long ,device=_lowerCAmelCase							)
																				lowerCamelCase__								=	torch.zeros(
																				    self.model_tester.batch_size ,dtype=torch.long ,device=_lowerCAmelCase							)
								return inputs_dict






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	MegatronBertModelTester(self							)
								lowerCamelCase__								=	ConfigTester(self ,config_class=_lowerCAmelCase ,hidden_size=37							)






		def       UpperCamelCase_       (   self							):
								self.config_tester.run_common_tests()






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_model(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_pretraining(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_question_answering(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_token_classification(*_lowerCAmelCase							)



def     A__      (					__lowerCAmelCase		:      Optional[Any]    ):
						return torch.tensor(
						    SCREAMING_SNAKE_CASE_ ,     dtype=torch.long ,     device=SCREAMING_SNAKE_CASE_ ,     )


UpperCamelCase     :   Any	      =		1E-4





@require_torch
@require_sentencepiece
@require_tokenizers
class 		UpperCamelCase__     (unittest.TestCase      ):




		'''simple docstring'''





		@slow
		@unittest.skip("""Model is not available."""							)
		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	"""nvidia/megatron-bert-uncased-345m"""
								if "MYDIR" in os.environ:
														lowerCamelCase__								=	os.path.join(os.environ["""MYDIR"""] ,_lowerCAmelCase							)
								lowerCamelCase__								=	MegatronBertModel.from_pretrained(_lowerCAmelCase							)
								model.to(_lowerCAmelCase							)
								model.half()
								lowerCamelCase__								=	_long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]							)
								with torch.no_grad():
														lowerCamelCase__								=	model(_lowerCAmelCase							)[0]
								lowerCamelCase__								=	torch.Size((1, 9, 10_24)							)
								self.assertEqual(output.shape ,_lowerCAmelCase							)

								lowerCamelCase__								=	[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
								for ii in range(3							):
														for jj in range(3							):
																				lowerCamelCase__								=	output[0, ii, jj]
																				lowerCamelCase__								=	expected[3 * ii + jj]
																				lowerCamelCase__								=	"""ii={} jj={} a={} b={}""".format(_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							)
																				self.assertTrue(math.isclose(_lowerCAmelCase ,_lowerCAmelCase ,rel_tol=_lowerCAmelCase ,abs_tol=_lowerCAmelCase							) ,msg=_lowerCAmelCase							)








'''simple docstring'''


from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


_SCREAMING_SNAKE_CASE									=				logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE									=				{
    "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json",
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}

class  lowerCAmelCase_		(		__magic_name__		):
 __lowerCamelCase    :				Union[str, Any]								=       "blenderbot-small"
 __lowerCamelCase    :				Optional[Any]								=       ["past_key_values"]
 __lowerCamelCase    :				str								=       {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




 def __init__(						self      ,  _lowerCAmelCase=50265      ,  _lowerCAmelCase=512      ,  _lowerCAmelCase=8      ,  _lowerCAmelCase=2048      ,  _lowerCAmelCase=16      ,  _lowerCAmelCase=8      ,  _lowerCAmelCase=2048      ,  _lowerCAmelCase=16      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase="gelu"      ,  _lowerCAmelCase=512      ,  _lowerCAmelCase=0.1      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=0.02      ,  _lowerCAmelCase=1      ,  _lowerCAmelCase=False      ,  _lowerCAmelCase=0      ,  _lowerCAmelCase=1      ,  _lowerCAmelCase=2      ,  _lowerCAmelCase=2      ,  **_lowerCAmelCase      ,  )     ->				Dict:
     _lowerCAmelCase						      =   vocab_size
     _lowerCAmelCase						      =   max_position_embeddings
     _lowerCAmelCase						      =   d_model
     _lowerCAmelCase						      =   encoder_ffn_dim
     _lowerCAmelCase						      =   encoder_layers
     _lowerCAmelCase						      =   encoder_attention_heads
     _lowerCAmelCase						      =   decoder_ffn_dim
     _lowerCAmelCase						      =   decoder_layers
     _lowerCAmelCase						      =   decoder_attention_heads
     _lowerCAmelCase						      =   dropout
     _lowerCAmelCase						      =   attention_dropout
     _lowerCAmelCase						      =   activation_dropout
     _lowerCAmelCase						      =   activation_function
     _lowerCAmelCase						      =   init_std
     _lowerCAmelCase						      =   encoder_layerdrop
     _lowerCAmelCase						      =   decoder_layerdrop
     _lowerCAmelCase						      =   use_cache
     _lowerCAmelCase						      =   encoder_layers
     _lowerCAmelCase						      =   scale_embedding  # scale factor will be sqrt(d_model) if True

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







class  lowerCAmelCase_		(		__magic_name__		):




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

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

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




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = False      ,  _lowerCAmelCase = None      ,  )     ->				Mapping[str, Any]:
     _lowerCAmelCase						      =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)

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

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

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

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

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




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = False      ,  _lowerCAmelCase = None      ,  )     ->				Mapping[str, Any]:
     _lowerCAmelCase						      =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)

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

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




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = False      ,  _lowerCAmelCase = 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
     _lowerCAmelCase						      =   compute_effective_axis_dimension(
         _lowerCAmelCase      ,  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
     _lowerCAmelCase						      =   tokenizer.num_special_tokens_to_add(_lowerCAmelCase		)
     _lowerCAmelCase						      =   compute_effective_axis_dimension(
         _lowerCAmelCase      ,  fixed_dimension=OnnxConfig.default_fixed_sequence      ,  num_token_to_add=_lowerCAmelCase		)

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




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = -1      ,  _lowerCAmelCase = False      ,  _lowerCAmelCase = None      ,  )     ->				Mapping[str, Any]:
     if self.task in ["default", "seq2seq-lm"]:
         _lowerCAmelCase						      =   self._generate_dummy_inputs_for_default_and_seqaseq_lm(
             _lowerCAmelCase      ,  batch_size=_lowerCAmelCase      ,  seq_length=_lowerCAmelCase      ,  is_pair=_lowerCAmelCase      ,  framework=_lowerCAmelCase		)

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

     return common_inputs





 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				Dict:
     if self.task in ["default", "seq2seq-lm"]:
         _lowerCAmelCase						      =   super()._flatten_past_key_values_(_lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)
     else:
         _lowerCAmelCase						      =   super(_lowerCAmelCase      ,  self		)._flatten_past_key_values_(
             _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)


