






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 		SCREAMING_SNAKE_CASE__	:

  '''simple docstring'''







  def __init__(      self		:			List[str]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			List[Any]=13       ,				lowerCamelCase		:			List[Any]=7       ,				lowerCamelCase		:			int=True       ,				lowerCamelCase		:			int=True       ,				lowerCamelCase		:			Dict=True       ,				lowerCamelCase		:			Optional[int]=True       ,				lowerCamelCase		:			Optional[Any]=99       ,				lowerCamelCase		:			Tuple=64       ,				lowerCamelCase		:			List[Any]=32       ,				lowerCamelCase		:			Tuple=5       ,				lowerCamelCase		:			Tuple=4       ,				lowerCamelCase		:			Union[str, Any]=37       ,				lowerCamelCase		:			List[Any]="gelu"       ,				lowerCamelCase		:			int=0.1       ,				lowerCamelCase		:			int=0.1       ,				lowerCamelCase		:			List[Any]=512       ,				lowerCamelCase		:			Tuple=16       ,				lowerCamelCase		:			Optional[Any]=2       ,				lowerCamelCase		:			int=0.02       ,				lowerCamelCase		:			Any=3       ,				lowerCamelCase		:			Dict=4       ,				lowerCamelCase		:			Tuple=None       ,				)						->	int:


         """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  lowerCamelCase       (      self		:			Union[str, Any]							)						->	List[Any]:


         """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  lowerCamelCase       (      self		:			Any							)						->	Any:


         """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=lowerCamelCase       ,				initializer_range=self.initializer_range       ,				)


  def  lowerCamelCase       (      self		:			List[Any]       ,				lowerCamelCase		:			Optional[int]       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			Optional[int]       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			Optional[int]							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertModel(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase							)
         _UpperCAmelCase							=    model(lowerCamelCase       ,				token_type_ids=lowerCamelCase							)
         _UpperCAmelCase							=    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  lowerCamelCase       (      self		:			List[Any]       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Any       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			List[Any]							)						->	int:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForMaskedLM(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    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  lowerCamelCase       (      self		:			Tuple       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			Tuple       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Tuple       ,				lowerCamelCase		:			Union[str, Any]							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForCausalLM(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    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  lowerCamelCase       (      self		:			str       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			str       ,				lowerCamelCase		:			Optional[int]							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForNextSentencePrediction(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(
             lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase       ,				)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, 2)							)


  def  lowerCamelCase       (      self		:			int       ,				lowerCamelCase		:			Any       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Optional[int]       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			List[Any]							)						->	Any:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForPreTraining(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    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  lowerCamelCase       (      self		:			Any       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			Any       ,				lowerCamelCase		:			int							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertForQuestionAnswering(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    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  lowerCamelCase       (      self		:			Union[str, Any]       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			Tuple       ,				lowerCamelCase		:			Optional[int]       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			List[Any]							)						->	Tuple:


         """simple docstring"""
         _UpperCAmelCase							=    self.num_labels
         _UpperCAmelCase							=    MegatronBertForSequenceClassification(lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    model(lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.num_labels)							)


  def  lowerCamelCase       (      self		:			int       ,				lowerCamelCase		:			Optional[int]       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			Optional[int]       ,				lowerCamelCase		:			str							)						->	int:


         """simple docstring"""
         _UpperCAmelCase							=    self.num_labels
         _UpperCAmelCase							=    MegatronBertForTokenClassification(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         model.eval()
         _UpperCAmelCase							=    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  lowerCamelCase       (      self		:			str       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Union[str, Any]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			List[Any]       ,				lowerCamelCase		:			List[str]       ,				lowerCamelCase		:			Dict       ,				lowerCamelCase		:			Tuple							)						->	List[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.num_choices
         _UpperCAmelCase							=    MegatronBertForMultipleChoice(config=lowerCamelCase							)
         model.to(lowerCamelCase							)
         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(
             lowerCamelCase       ,				attention_mask=lowerCamelCase       ,				token_type_ids=lowerCamelCase       ,				labels=lowerCamelCase       ,				)
         self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.num_choices)							)





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


         """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 		SCREAMING_SNAKE_CASE__	(  UpperCAmelCase					,      UpperCAmelCase					,      unittest.TestCase     ):

  '''simple docstring'''







  _lowerCamelCase    =  (
      (
          MegatronBertModel,
          MegatronBertForMaskedLM,
          MegatronBertForCausalLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  _lowerCamelCase    =  (
      {
          '''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 {}
  )
  _lowerCamelCase    =  True
  # test_resize_embeddings = False
  _lowerCamelCase    =  False


  def  lowerCamelCase       (      self		:			Any       ,				lowerCamelCase		:			Optional[Any]       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Optional[int]=False							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    super()._prepare_for_class(lowerCamelCase       ,				lowerCamelCase       ,				return_labels=lowerCamelCase							)

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


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


         """simple docstring"""
         _UpperCAmelCase							=    MegatronBertModelTester(self							)
         _UpperCAmelCase							=    ConfigTester(self       ,				config_class=lowerCamelCase       ,				hidden_size=37							)


  def  lowerCamelCase       (      self		:			List[str]							)						->	Any:


         """simple docstring"""
         self.config_tester.run_common_tests()


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


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase							)


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


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Dict							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			List[Any]							)						->	Optional[Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Union[str, Any]							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			Tuple							)						->	Dict:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase							)


  def  lowerCamelCase       (      self		:			List[str]							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase							)





  def  lowerCamelCase       (      self		:			Dict							)						->	List[str]:


         """simple docstring"""
         _UpperCAmelCase							=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase							)




def       _SCREAMING_SNAKE_CASE (							__snake_case )				->						Dict:
       return torch.tensor(
           __snake_case   ,      dtype=torch.long   ,      device=__snake_case   ,      )


__a:     int   					=					1E-4





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

  '''simple docstring'''







  @slow
  @unittest.skip("""Model is not available."""							)
  def  lowerCamelCase       (      self		:			int							)						->	Union[str, Any]:


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

         _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(lowerCamelCase       ,				lowerCamelCase       ,				lowerCamelCase       ,				lowerCamelCase							)
                       self.assertTrue(math.isclose(lowerCamelCase       ,				lowerCamelCase       ,				rel_tol=lowerCamelCase       ,				abs_tol=lowerCamelCase							)       ,				msg=lowerCamelCase							)



import inspect
import unittest






class      UpperCamelCase(	unittest.TestCase ):


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Dict    )      ->						List[Any]:




											'''simple docstring'''

											try:
																	import diffusers  # noqa: F401
											except ImportError:
																	assert False






					def 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[int]    )      ->						Any:




											'''simple docstring'''

											import diffusers
											from diffusers.dependency_versions_table import deps

											__snake_case				=		inspect.getmembers(SCREAMING_SNAKE_CASE  ,     inspect.isclass    )

											for cls_name, cls_module in all_classes:
																	if "dummy_" in cls_module.__module__:
																							for backend in cls_module._backends:
																													if backend == "k_diffusion":
																																			__snake_case				=		"k-diffusion"
																													elif backend == "invisible_watermark":
																																			__snake_case				=		"invisible-watermark"
																													assert backend in deps, f'''{backend} is not in the deps table!'''





