


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 		__a				:
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=13      ,   _SCREAMING_SNAKE_CASE=7      ,   _SCREAMING_SNAKE_CASE=True      ,   _SCREAMING_SNAKE_CASE=True      ,   _SCREAMING_SNAKE_CASE=True      ,   _SCREAMING_SNAKE_CASE=True      ,   _SCREAMING_SNAKE_CASE=99      ,   _SCREAMING_SNAKE_CASE=64      ,   _SCREAMING_SNAKE_CASE=32      ,   _SCREAMING_SNAKE_CASE=5      ,   _SCREAMING_SNAKE_CASE=4      ,   _SCREAMING_SNAKE_CASE=37      ,   _SCREAMING_SNAKE_CASE="gelu"      ,   _SCREAMING_SNAKE_CASE=0.1      ,   _SCREAMING_SNAKE_CASE=0.1      ,   _SCREAMING_SNAKE_CASE=512      ,   _SCREAMING_SNAKE_CASE=16      ,   _SCREAMING_SNAKE_CASE=2      ,   _SCREAMING_SNAKE_CASE=0.02      ,   _SCREAMING_SNAKE_CASE=3      ,   _SCREAMING_SNAKE_CASE=4      ,   _SCREAMING_SNAKE_CASE=None      ,   )			->	Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      parent
									_UpperCAmelCase					=      batch_size
									_UpperCAmelCase					=      seq_length
									_UpperCAmelCase					=      is_training
									_UpperCAmelCase					=      use_input_mask
									_UpperCAmelCase					=      use_token_type_ids
									_UpperCAmelCase					=      use_labels
									_UpperCAmelCase					=      vocab_size
									_UpperCAmelCase					=      hidden_size
									_UpperCAmelCase					=      embedding_size
									_UpperCAmelCase					=      num_hidden_layers
									_UpperCAmelCase					=      num_attention_heads
									_UpperCAmelCase					=      intermediate_size
									_UpperCAmelCase					=      hidden_act
									_UpperCAmelCase					=      hidden_dropout_prob
									_UpperCAmelCase					=      attention_probs_dropout_prob
									_UpperCAmelCase					=      max_position_embeddings
									_UpperCAmelCase					=      type_vocab_size
									_UpperCAmelCase					=      type_sequence_label_size
									_UpperCAmelCase					=      initializer_range
									_UpperCAmelCase					=      num_labels
									_UpperCAmelCase					=      num_choices
									_UpperCAmelCase					=      scope
		def      UpperCAmelCase__    (      self						)			->	List[str]:




									"""simple docstring"""





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

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

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

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

									_UpperCAmelCase					=      self.get_config()

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
		def      UpperCAmelCase__    (      self						)			->	List[str]:




									"""simple docstring"""





									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=_SCREAMING_SNAKE_CASE      ,   initializer_range=self.initializer_range      ,   )
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      MegatronBertModel(config=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE						)

									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      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      MegatronBertForMaskedLM(config=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE      ,   labels=_SCREAMING_SNAKE_CASE						)
									self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size)						)
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=      MegatronBertForCausalLM(config=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE      ,   labels=_SCREAMING_SNAKE_CASE						)
									self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.vocab_size)						)
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=      MegatronBertForNextSentencePrediction(config=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(
									    _SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE      ,   labels=_SCREAMING_SNAKE_CASE      ,   )
									self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, 2)						)
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      MegatronBertForPreTraining(config=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(
									    _SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE      ,   labels=_SCREAMING_SNAKE_CASE      ,   next_sentence_label=_SCREAMING_SNAKE_CASE      ,   )
									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      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      MegatronBertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(
									    _SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE      ,   start_positions=_SCREAMING_SNAKE_CASE      ,   end_positions=_SCREAMING_SNAKE_CASE      ,   )
									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      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.num_labels
									_UpperCAmelCase					=      MegatronBertForSequenceClassification(_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE      ,   labels=_SCREAMING_SNAKE_CASE						)
									self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.num_labels)						)
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.num_labels
									_UpperCAmelCase					=      MegatronBertForTokenClassification(config=_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.eval()
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   token_type_ids=_SCREAMING_SNAKE_CASE      ,   labels=_SCREAMING_SNAKE_CASE						)
									self.parent.assertEqual(result.logits.shape      ,   (self.batch_size, self.seq_length, self.num_labels)						)
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Union[str, Any]:




									"""simple docstring"""





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






		def      UpperCAmelCase__    (      self						)			->	Dict:




									"""simple docstring"""





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



@require_torch
class 		__a				(    UpperCAmelCase   ,				UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  Optional[int]           =     (
		    (
		        MegatronBertModel,
		        MegatronBertForMaskedLM,
		        MegatronBertForCausalLM,
		        MegatronBertForMultipleChoice,
		        MegatronBertForNextSentencePrediction,
		        MegatronBertForPreTraining,
		        MegatronBertForQuestionAnswering,
		        MegatronBertForSequenceClassification,
		        MegatronBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		_a     :  Tuple           =     (
		    {
		        '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 {}
		)
		_a     :  Tuple           =     True
		# test_resize_embeddings = False
		_a     :  Optional[int]           =     False
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=False						)			->	List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=      super()._prepare_for_class(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   return_labels=_SCREAMING_SNAKE_CASE						)

									if return_labels:
																if model_class in get_values(_SCREAMING_SNAKE_CASE						):
																							_UpperCAmelCase					=      torch.zeros(
																							    (self.model_tester.batch_size, self.model_tester.seq_length)      ,   dtype=torch.long      ,   device=_SCREAMING_SNAKE_CASE						)
																							_UpperCAmelCase					=      torch.zeros(
																							    self.model_tester.batch_size      ,   dtype=torch.long      ,   device=_SCREAMING_SNAKE_CASE						)
									return inputs_dict
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      MegatronBertModelTester(self						)
									_UpperCAmelCase					=      ConfigTester(self      ,   config_class=_SCREAMING_SNAKE_CASE      ,   hidden_size=37						)
		def      UpperCAmelCase__    (      self						)			->	Union[str, Any]:




									"""simple docstring"""





									self.config_tester.run_common_tests()
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_model(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_pretraining(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_question_answering(*_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_SCREAMING_SNAKE_CASE						)






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




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_token_classification(*_SCREAMING_SNAKE_CASE						)



def 	lowerCAmelCase__			(      a__:		Dict	)  ->		Tuple:


							'''simple docstring'''





							return torch.tensor(
							    a__      ,  dtype=torch.long      ,  device=a__      ,  )


lowerCAmelCase__  :Union[str, Any]             =							1e-4







@require_torch
@require_sentencepiece
@require_tokenizers
class 		__a				(    unittest.TestCase      ):
		@slow
		@unittest.skip('Model is not available.'						)
		def      UpperCAmelCase__    (      self						)			->	Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=      'nvidia/megatron-bert-uncased-345m'
									if "MYDIR" in os.environ:
																_UpperCAmelCase					=      os.path.join(os.environ['MYDIR']      ,   _SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      MegatronBertModel.from_pretrained(_SCREAMING_SNAKE_CASE						)
									model.to(_SCREAMING_SNAKE_CASE						)
									model.half()
									_UpperCAmelCase					=      _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]						)
									with torch.no_grad():
																_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE						)[0]
									_UpperCAmelCase					=      torch.Size((1, 9, 1024)						)
									self.assertEqual(output.shape      ,   _SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase					=      [-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						):
																							_UpperCAmelCase					=      output[0, ii, jj]
																							_UpperCAmelCase					=      expected[3 * ii + jj]
																							_UpperCAmelCase					=      'ii={} jj={} a={} b={}'.format(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
																							self.assertTrue(math.isclose(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   rel_tol=_SCREAMING_SNAKE_CASE      ,   abs_tol=_SCREAMING_SNAKE_CASE						)      ,   msg=_SCREAMING_SNAKE_CASE						)




from __future__ import annotations

import unittest

from transformers import BlenderbotConfig, BlenderbotTokenizer, is_tf_available
from transformers.testing_utils import require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers import TFAutoModelForSeqaSeqLM, TFBlenderbotForConditionalGeneration, TFBlenderbotModel

@require_tf
class 		__a				:
		_a     :  Dict           =     BlenderbotConfig
		_a     :  Dict           =     {}
		_a     :  Union[str, Any]           =     'gelu'
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=13      ,   _SCREAMING_SNAKE_CASE=7      ,   _SCREAMING_SNAKE_CASE=True      ,   _SCREAMING_SNAKE_CASE=False      ,   _SCREAMING_SNAKE_CASE=99      ,   _SCREAMING_SNAKE_CASE=32      ,   _SCREAMING_SNAKE_CASE=2      ,   _SCREAMING_SNAKE_CASE=4      ,   _SCREAMING_SNAKE_CASE=37      ,   _SCREAMING_SNAKE_CASE=0.1      ,   _SCREAMING_SNAKE_CASE=0.1      ,   _SCREAMING_SNAKE_CASE=20      ,   _SCREAMING_SNAKE_CASE=2      ,   _SCREAMING_SNAKE_CASE=1      ,   _SCREAMING_SNAKE_CASE=0      ,   )			->	Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase					=      parent
									_UpperCAmelCase					=      batch_size
									_UpperCAmelCase					=      seq_length
									_UpperCAmelCase					=      is_training
									_UpperCAmelCase					=      use_labels
									_UpperCAmelCase					=      vocab_size
									_UpperCAmelCase					=      hidden_size
									_UpperCAmelCase					=      num_hidden_layers
									_UpperCAmelCase					=      num_attention_heads
									_UpperCAmelCase					=      intermediate_size

									_UpperCAmelCase					=      hidden_dropout_prob
									_UpperCAmelCase					=      attention_probs_dropout_prob
									_UpperCAmelCase					=      max_position_embeddings
									_UpperCAmelCase					=      eos_token_id
									_UpperCAmelCase					=      pad_token_id
									_UpperCAmelCase					=      bos_token_id
		def      UpperCAmelCase__    (      self						)			->	Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=      ids_tensor([self.batch_size, self.seq_length - 1]      ,   self.vocab_size						)
									_UpperCAmelCase					=      tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size						)      ,   1						)
									_UpperCAmelCase					=      tf.concat([input_ids, eos_tensor]      ,   axis=1						)

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

									_UpperCAmelCase					=      self.config_cls(
									    vocab_size=self.vocab_size      ,   d_model=self.hidden_size      ,   encoder_layers=self.num_hidden_layers      ,   decoder_layers=self.num_hidden_layers      ,   encoder_attention_heads=self.num_attention_heads      ,   decoder_attention_heads=self.num_attention_heads      ,   encoder_ffn_dim=self.intermediate_size      ,   decoder_ffn_dim=self.intermediate_size      ,   dropout=self.hidden_dropout_prob      ,   attention_dropout=self.attention_probs_dropout_prob      ,   max_position_embeddings=self.max_position_embeddings      ,   eos_token_ids=[2]      ,   bos_token_id=self.bos_token_id      ,   pad_token_id=self.pad_token_id      ,   decoder_start_token_id=self.pad_token_id      ,   **self.config_updates      ,   )
									_UpperCAmelCase					=      prepare_blenderbot_inputs_dict(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									return config, inputs_dict






		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      TFBlenderbotModel(config=_SCREAMING_SNAKE_CASE						).get_decoder()
									_UpperCAmelCase					=      inputs_dict['input_ids']

									_UpperCAmelCase					=      input_ids[:1, :]
									_UpperCAmelCase					=      inputs_dict['attention_mask'][:1, :]
									_UpperCAmelCase					=      inputs_dict['head_mask']
									_UpperCAmelCase					=      1

									# first forward pass
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   head_mask=_SCREAMING_SNAKE_CASE      ,   use_cache=_SCREAMING_SNAKE_CASE						)

									_UpperCAmelCase		,	_UpperCAmelCase					=      outputs.to_tuple()

									# create hypothetical next token and extent to next_input_ids
									_UpperCAmelCase					=      ids_tensor((self.batch_size, 3)      ,   config.vocab_size						)
									_UpperCAmelCase					=      tf.cast(ids_tensor((self.batch_size, 3)      ,   2						)      ,   tf.inta						)

									# append to next input_ids and
									_UpperCAmelCase					=      tf.concat([input_ids, next_tokens]      ,   axis=-1						)
									_UpperCAmelCase					=      tf.concat([attention_mask, next_attn_mask]      ,   axis=-1						)

									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE						)[0]
									_UpperCAmelCase					=      model(_SCREAMING_SNAKE_CASE      ,   attention_mask=_SCREAMING_SNAKE_CASE      ,   past_key_values=_SCREAMING_SNAKE_CASE						)[0]

									self.parent.assertEqual(next_tokens.shape[1]      ,   output_from_past.shape[1]						)

									# select random slice
									_UpperCAmelCase					=      int(ids_tensor((1,)      ,   output_from_past.shape[-1]						)						)
									_UpperCAmelCase					=      output_from_no_past[:, -3:, random_slice_idx]
									_UpperCAmelCase					=      output_from_past[:, :, random_slice_idx]

									# test that outputs are equal for slice
									tf.debugging.assert_near(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   rtol=1e-3						)



def 	lowerCAmelCase__			(      a__:		Dict      ,  a__:		Dict      ,  a__:		Any      ,  a__:		Any=None      ,  a__:		List[Any]=None      ,  a__:		Union[str, Any]=None      ,  a__:		Tuple=None      ,  a__:		Union[str, Any]=None      ,  )  ->		Any:


							'''simple docstring'''





							if attention_mask is None:
														_UpperCAmelCase					=      tf.cast(tf.math.not_equal(a__      ,  config.pad_token_id	)      ,  tf.inta	)
							if decoder_attention_mask is None:
														_UpperCAmelCase					=      tf.concat(
														    [
														        tf.ones(decoder_input_ids[:, :1].shape      ,  dtype=tf.inta	),
														        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]      ,  config.pad_token_id	)      ,  tf.inta	),
														    ]      ,  axis=-1      ,  )
							if head_mask is None:
														_UpperCAmelCase					=      tf.ones((config.encoder_layers, config.encoder_attention_heads)	)
							if decoder_head_mask is None:
														_UpperCAmelCase					=      tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
							if cross_attn_head_mask is None:
														_UpperCAmelCase					=      tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
							return {
							    "input_ids": input_ids,
							    "decoder_input_ids": decoder_input_ids,
							    "attention_mask": attention_mask,
							    "decoder_attention_mask": decoder_attention_mask,
							    "head_mask": head_mask,
							    "decoder_head_mask": decoder_head_mask,
							    "cross_attn_head_mask": cross_attn_head_mask,
							}

@require_tf
class 		__a				(    UpperCAmelCase   ,				UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  List[Any]           =     (TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
		_a     :  List[str]           =     (TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
		_a     :  List[str]           =     (
		    {
		        'conversational': TFBlenderbotForConditionalGeneration,
		        'feature-extraction': TFBlenderbotModel,
		        'summarization': TFBlenderbotForConditionalGeneration,
		        'text2text-generation': TFBlenderbotForConditionalGeneration,
		        'translation': TFBlenderbotForConditionalGeneration,
		    }
		    if is_tf_available()
		    else {}
		)
		_a     :  Dict           =     True
		_a     :  int           =     False
		_a     :  Union[str, Any]           =     False
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      TFBlenderbotModelTester(self						)
									_UpperCAmelCase					=      ConfigTester(self      ,   config_class=_SCREAMING_SNAKE_CASE						)
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									self.config_tester.run_common_tests()






		def      UpperCAmelCase__    (      self						)			->	List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=      self.model_tester.prepare_config_and_inputs_for_common()
									self.model_tester.check_decoder_model_past_large_inputs(*_SCREAMING_SNAKE_CASE						)



@require_tokenizers
@require_tf
class 		__a				(    unittest.TestCase      ):
		_a     :  int           =     ['My friends are cool but they eat too many carbs.']
		_a     :  List[Any]           =     'facebook/blenderbot-400M-distill'
		@cached_property
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									return BlenderbotTokenizer.from_pretrained(self.model_name						)
		@cached_property
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name						)
									return model






		@slow
		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									_UpperCAmelCase					=      self.tokenizer(self.src_text      ,   return_tensors='tf'						)
									_UpperCAmelCase					=      self.model.generate(
									    model_inputs.input_ids      ,   )
									_UpperCAmelCase					=      self.tokenizer.batch_decode(generated_ids.numpy()      ,   skip_special_tokens=_SCREAMING_SNAKE_CASE						)[0]
									assert (
									    generated_words
									    == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?"
									)
