"""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_      :


      def __init__(  self						,     a						,     a=1_3						,     a=7						,     a=True						,     a=True						,     a=True						,     a=True						,     a=9_9						,     a=6_4						,     a=3_2						,     a=5						,     a=4						,     a=3_7						,     a="gelu"						,     a=0.1						,     a=0.1						,     a=5_1_2						,     a=1_6						,     a=2						,     a=0.02						,     a=3						,     a=4						,     a=None						,     )			->   List[Any]:
             lowercase__     :						List[Any]         =      parent
             lowercase__     :						str         =      batch_size
             lowercase__     :						List[str]         =      seq_length
             lowercase__     :						Dict         =      is_training
             lowercase__     :						Dict         =      use_input_mask
             lowercase__     :						str         =      use_token_type_ids
             lowercase__     :						Tuple         =      use_labels
             lowercase__     :						int         =      vocab_size
             lowercase__     :						Optional[Any]         =      hidden_size
             lowercase__     :						Tuple         =      embedding_size
             lowercase__     :						str         =      num_hidden_layers
             lowercase__     :						Union[str, Any]         =      num_attention_heads
             lowercase__     :						Optional[Any]         =      intermediate_size
             lowercase__     :						Tuple         =      hidden_act
             lowercase__     :						Union[str, Any]         =      hidden_dropout_prob
             lowercase__     :						Union[str, Any]         =      attention_probs_dropout_prob
             lowercase__     :						Union[str, Any]         =      max_position_embeddings
             lowercase__     :						Tuple         =      type_vocab_size
             lowercase__     :						str         =      type_sequence_label_size
             lowercase__     :						int         =      initializer_range
             lowercase__     :						int         =      num_labels
             lowercase__     :						str         =      num_choices
             lowercase__     :						Tuple         =      scope


      def 						_UpperCAmelCase			(  self				)			->   Tuple:
             lowercase__     :						Tuple         =      ids_tensor([self.batch_size, self.seq_length]						,     self.vocab_size				)

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

             lowercase__     :						Optional[Any]         =      None
             if self.use_token_type_ids:
                    lowercase__     :						Union[str, Any]         =      ids_tensor([self.batch_size, self.seq_length]						,     self.type_vocab_size				)

             lowercase__     :						Optional[Any]         =      None
             lowercase__     :						Union[str, Any]         =      None
             lowercase__     :						List[str]         =      None
             if self.use_labels:
                    lowercase__     :						Any         =      ids_tensor([self.batch_size]						,     self.type_sequence_label_size				)
                    lowercase__     :						List[str]         =      ids_tensor([self.batch_size, self.seq_length]						,     self.num_labels				)
                    lowercase__     :						Union[str, Any]         =      ids_tensor([self.batch_size]						,     self.num_choices				)

             lowercase__     :						Optional[Any]         =      self.get_config()

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


      def 						_UpperCAmelCase			(  self				)			->   List[Any]:
             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=lowercase__						,     initializer_range=self.initializer_range						,     )


      def 						_UpperCAmelCase			(  self						,     a						,     a						,     a						,     a						,     a						,     a						,     a				)			->   int:
             lowercase__     :						str         =      MegatronBertModel(config=lowercase__				)
             model.to(lowercase__				)
             model.eval()
             lowercase__     :						Dict         =      model(lowercase__						,     attention_mask=lowercase__						,     token_type_ids=lowercase__				)
             lowercase__     :						Optional[int]         =      model(lowercase__						,     token_type_ids=lowercase__				)
             lowercase__     :						List[str]         =      model(lowercase__				)

             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						,     a						,     a						,     a						,     a						,     a						,     a						,     a				)			->   List[Any]:
             lowercase__     :						Union[str, Any]         =      MegatronBertForMaskedLM(config=lowercase__				)
             model.to(lowercase__				)
             model.eval()
             lowercase__     :						Any         =      model(lowercase__						,     attention_mask=lowercase__						,     token_type_ids=lowercase__						,     labels=lowercase__				)
             self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.seq_length, self.vocab_size)				)


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


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


      def 						_UpperCAmelCase			(  self						,     a						,     a						,     a						,     a						,     a						,     a						,     a				)			->   Tuple:
             lowercase__     :						Tuple         =      MegatronBertForPreTraining(config=lowercase__				)
             model.to(lowercase__				)
             model.eval()
             lowercase__     :						Optional[Any]         =      model(
                 lowercase__						,     attention_mask=lowercase__						,     token_type_ids=lowercase__						,     labels=lowercase__						,     next_sentence_label=lowercase__						,     )
             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						,     a						,     a						,     a						,     a						,     a						,     a						,     a				)			->   str:
             lowercase__     :						Optional[int]         =      MegatronBertForQuestionAnswering(config=lowercase__				)
             model.to(lowercase__				)
             model.eval()
             lowercase__     :						List[Any]         =      model(
                 lowercase__						,     attention_mask=lowercase__						,     token_type_ids=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						,     a						,     a						,     a						,     a						,     a						,     a						,     a				)			->   Any:
             lowercase__     :						Union[str, Any]         =      self.num_labels
             lowercase__     :						List[Any]         =      MegatronBertForSequenceClassification(lowercase__				)
             model.to(lowercase__				)
             model.eval()
             lowercase__     :						Tuple         =      model(lowercase__						,     attention_mask=lowercase__						,     token_type_ids=lowercase__						,     labels=lowercase__				)
             self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_labels)				)


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


      def 						_UpperCAmelCase			(  self						,     a						,     a						,     a						,     a						,     a						,     a						,     a				)			->   Union[str, Any]:
             lowercase__     :						str         =      self.num_choices
             lowercase__     :						str         =      MegatronBertForMultipleChoice(config=lowercase__				)
             model.to(lowercase__				)
             model.eval()
             lowercase__     :						Any         =      input_ids.unsqueeze(1				).expand(-1						,     self.num_choices						,     -1				).contiguous()
             lowercase__     :						Tuple         =      token_type_ids.unsqueeze(1				).expand(-1						,     self.num_choices						,     -1				).contiguous()
             lowercase__     :						Dict         =      input_mask.unsqueeze(1				).expand(-1						,     self.num_choices						,     -1				).contiguous()
             lowercase__     :						Tuple         =      model(
                 lowercase__						,     attention_mask=lowercase__						,     token_type_ids=lowercase__						,     labels=lowercase__						,     )
             self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_choices)				)


      def 						_UpperCAmelCase			(  self				)			->   List[str]:
             lowercase__     :						int         =      self.prepare_config_and_inputs()
             (
                 (
                 lowercase__		
             )       , (
                 lowercase__		
             )       , (
                 lowercase__		
             )       , (
                 lowercase__		
             )       , (
                 lowercase__		
             )       , (
                 lowercase__		
             )       , (
                 lowercase__		
             )       , 		
             )     :						Any         =      config_and_inputs
             lowercase__     :						Optional[int]         =      {'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):
      lowerCamelCase__						:       Dict					    =    (
          (
              MegatronBertModel,
              MegatronBertForMaskedLM,
              MegatronBertForCausalLM,
              MegatronBertForMultipleChoice,
              MegatronBertForNextSentencePrediction,
              MegatronBertForPreTraining,
              MegatronBertForQuestionAnswering,
              MegatronBertForSequenceClassification,
              MegatronBertForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      lowerCamelCase__						:       Any					    =    (
          {
              "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__						:       Optional[Any]					    =    True
      # test_resize_embeddings = False
      lowerCamelCase__						:       Any					    =    False


      def 						_UpperCAmelCase			(  self						,     a						,     a						,     a=False				)			->   Union[str, Any]:
             lowercase__     :						List[Any]         =      super()._prepare_for_class(lowercase__						,     lowercase__						,     return_labels=lowercase__				)

             if return_labels:
                    if model_class in get_values(lowercase__				):
                           lowercase__     :						List[str]         =      torch.zeros(
                               (self.model_tester.batch_size, self.model_tester.seq_length)						,     dtype=torch.long						,     device=lowercase__				)
                           lowercase__     :						Tuple         =      torch.zeros(
                               self.model_tester.batch_size						,     dtype=torch.long						,     device=lowercase__				)
             return inputs_dict


      def 						_UpperCAmelCase			(  self				)			->   int:
             lowercase__     :						Union[str, Any]         =      MegatronBertModelTester(self				)
             lowercase__     :						Dict         =      ConfigTester(self						,     config_class=lowercase__						,     hidden_size=3_7				)


      def 						_UpperCAmelCase			(  self				)			->   Tuple:
             self.config_tester.run_common_tests()


      def 						_UpperCAmelCase			(  self				)			->   str:
             lowercase__     :						Dict         =      self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_model(*lowercase__				)


      def 						_UpperCAmelCase			(  self				)			->   Union[str, Any]:
             lowercase__     :						Any         =      self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase__				)


      def 						_UpperCAmelCase			(  self				)			->   str:
             lowercase__     :						Optional[int]         =      self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase__				)


      def 						_UpperCAmelCase			(  self				)			->   List[Any]:
             lowercase__     :						str         =      self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase__				)


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


      def 						_UpperCAmelCase			(  self				)			->   Dict:
             lowercase__     :						Any         =      self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase__				)


      def 						_UpperCAmelCase			(  self				)			->   Optional[int]:
             lowercase__     :						int         =      self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase__				)


      def 						_UpperCAmelCase			(  self				)			->   Union[str, Any]:
             lowercase__     :						List[Any]         =      self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase__				)





def 							a_	( _lowerCAmelCase    :					Any	):





       '''simple docstring'''

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


_UpperCamelCase  :		Optional[int]           =       1e-4






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


      @slow
      @unittest.skip('Model is not available.'				)
      def 						_UpperCAmelCase			(  self				)			->   Union[str, Any]:
             lowercase__     :						Tuple         =      'nvidia/megatron-bert-uncased-345m'
             if "MYDIR" in os.environ:
                    lowercase__     :						Optional[Any]         =      os.path.join(os.environ['MYDIR']						,     lowercase__				)
             lowercase__     :						Any         =      MegatronBertModel.from_pretrained(lowercase__				)
             model.to(lowercase__				)
             model.half()
             lowercase__     :						Optional[int]         =      _long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]				)
             with torch.no_grad():
                    lowercase__     :						Optional[Any]         =      model(lowercase__				)[0]
             lowercase__     :						List[str]         =      torch.Size((1, 9, 1_0_2_4)				)
             self.assertEqual(output.shape						,     lowercase__				)

             lowercase__     :						int         =      [-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
             for ii in range(3				):
                    for jj in range(3				):
                           lowercase__     :						Any         =      output[0, ii, jj]
                           lowercase__     :						List[str]         =      expected[3 * ii + jj]
                           lowercase__     :						str         =      'ii={} jj={} a={} b={}'.format(lowercase__						,     lowercase__						,     lowercase__						,     lowercase__				)
                           self.assertTrue(math.isclose(lowercase__						,     lowercase__						,     rel_tol=lowercase__						,     abs_tol=lowercase__				)						,     msg=lowercase__				)




import os
import re
import shutil
import sys
import tempfile
import unittest

import black


_lowerCamelCase      :					Union[str, Any]					    =      os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, '''utils'''))

import check_copies  # noqa: E402


# This is the reference code that will be used in the tests.
# If DDPMSchedulerOutput is changed in scheduling_ddpm.py, this code needs to be manually updated.
_lowerCamelCase      :					str					    =      '''    \"""
    Output class for the scheduler\'s step function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            The predicted denoised sample (x_{0}) based on the model output from the current timestep.
            `pred_original_sample` can be used to preview progress or for guidance.
    \"""

    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None
'''





class        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =tempfile.mkdtemp()
             os.makedirs(os.path.join(self.diffusers_dir ,   '''schedulers/'''     )     )
             __A    =self.diffusers_dir
             shutil.copy(
                 os.path.join(lowercase__ ,   '''src/diffusers/schedulers/scheduling_ddpm.py'''     ) ,   os.path.join(self.diffusers_dir ,   '''schedulers/scheduling_ddpm.py'''     ) ,   )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ='''src/diffusers'''
             shutil.rmtree(self.diffusers_dir     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__=None     ):







             '''simple docstring'''



             __A    =comment + f'''\nclass {class_name}(nn.Module):\n''' + class_code
             if overwrite_result is not None:
                   __A    =comment + f'''\nclass {class_name}(nn.Module):\n''' + overwrite_result
             __A    =black.Mode(target_versions={black.TargetVersion.PYaa} ,   line_length=1_1_9     )
             __A    =black.format_str(lowercase__ ,   mode=lowercase__     )
             __A    =os.path.join(self.diffusers_dir ,   '''new_code.py'''     )
             with open(lowercase__ ,   '''w''' ,   newline='''\n'''     ) as f:
                   f.write(lowercase__     )
             if overwrite_result is None:
                   self.assertTrue(len(check_copies.is_copy_consistent(lowercase__     )     ) == 0     )
             else:
                   check_copies.is_copy_consistent(f.name ,   overwrite=lowercase__     )
                   with open(lowercase__ ,   '''r'''     ) as f:
                         self.assertTrue(f.read() ,   lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =check_copies.find_code_in_diffusers('''schedulers.scheduling_ddpm.DDPMSchedulerOutput'''     )
             self.assertEqual(lowercase__ ,   lowercase__     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput''' ,   '''DDPMSchedulerOutput''' ,   REFERENCE_CODE + '''\n''' ,   )

             # With no empty line at the end
             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput''' ,   '''DDPMSchedulerOutput''' ,   lowercase__ ,   )

             # Copy consistency with rename
             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test''' ,   '''TestSchedulerOutput''' ,   re.sub('''DDPM''' ,   '''Test''' ,   lowercase__     ) ,   )

             # Copy consistency with a really long name
             __A    ='''TestClassWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReason'''
             self.check_copy_consistency(
                 f'''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->{long_class_name}''' ,   f'''{long_class_name}SchedulerOutput''' ,   re.sub('''Bert''' ,   lowercase__ ,   lowercase__     ) ,   )

             # Copy consistency with overwrite
             self.check_copy_consistency(
                 '''# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test''' ,   '''TestSchedulerOutput''' ,   lowercase__ ,   overwrite_result=re.sub('''DDPM''' ,   '''Test''' ,   lowercase__     ) ,   )


