



import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      NezhaForMaskedLM,
      NezhaForMultipleChoice,
      NezhaForNextSentencePrediction,
      NezhaForPreTraining,
      NezhaForQuestionAnswering,
      NezhaForSequenceClassification,
      NezhaForTokenClassification,
      NezhaModel,
  )
  from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST

class        snake_case__    :







      """simple docstring"""




      def __init__(       self ,       __lowercase ,       __lowercase=1_3 ,       __lowercase=7 ,       __lowercase=True ,       __lowercase=True ,       __lowercase=True ,       __lowercase=True ,       __lowercase=9_9 ,       __lowercase=3_2 ,       __lowercase=5 ,       __lowercase=4 ,       __lowercase=3_7 ,       __lowercase="gelu" ,       __lowercase=0.1 ,       __lowercase=0.1 ,       __lowercase=1_2_8 ,       __lowercase=3_2 ,       __lowercase=1_6 ,       __lowercase=2 ,       __lowercase=0.0_2 ,       __lowercase=3 ,       __lowercase=4 ,       __lowercase=None ,       ) ->				Dict:
         """simple docstring"""




         a__       :						List[str]							=  parent
         a__       :						Tuple							=  batch_size
         a__       :						List[Any]							=  seq_length
         a__       :						Tuple							=  is_training
         a__       :						Dict							=  use_input_mask
         a__       :						str							=  use_token_type_ids
         a__       :						List[Any]							=  use_labels
         a__       :						str							=  vocab_size
         a__       :						int							=  hidden_size
         a__       :						List[Any]							=  num_hidden_layers
         a__       :						str							=  num_attention_heads
         a__       :						Optional[int]							=  intermediate_size
         a__       :						Optional[Any]							=  hidden_act
         a__       :						Optional[Any]							=  hidden_dropout_prob
         a__       :						List[str]							=  attention_probs_dropout_prob
         a__       :						List[Any]							=  max_position_embeddings
         a__       :						List[str]							=  type_vocab_size
         a__       :						Any							=  type_sequence_label_size
         a__       :						List[str]							=  initializer_range
         a__       :						Any							=  num_labels
         a__       :						str							=  num_choices
         a__       :						int							=  scope

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[Any]:
         """simple docstring"""




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

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

         a__       :						List[str]							=  None
         if self.use_token_type_ids:
            a__       :						int							=  ids_tensor([self.batch_size, self.seq_length] ,       self.type_vocab_size  )

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

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

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

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Union[str, Any]:
         """simple docstring"""




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

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[int]:
         """simple docstring"""




         (
             a__						
         )       :						Optional[int]							=  self.prepare_config_and_inputs()

         a__       :						Dict							=  True
         a__       :						Optional[int]							=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
         a__       :						Any							=  ids_tensor([self.batch_size, self.seq_length] ,       vocab_size=2  )

         return (
             config,
             input_ids,
             token_type_ids,
             input_mask,
             sequence_labels,
             token_labels,
             choice_labels,
             encoder_hidden_states,
             encoder_attention_mask,
         )

      def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				List[str]:
         """simple docstring"""




         a__       :						Any							=  NezhaModel(config=__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						str							=  model(__lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase  )
         a__       :						Optional[Any]							=  model(__lowercase ,       token_type_ids=__lowercase  )
         a__       :						Optional[Any]							=  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 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       ) ->				str:
         """simple docstring"""




         a__       :						Any							=  True
         a__       :						Optional[Any]							=  NezhaModel(__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						Union[str, Any]							=  model(
             __lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       encoder_hidden_states=__lowercase ,       encoder_attention_mask=__lowercase ,       )
         a__       :						str							=  model(
             __lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       encoder_hidden_states=__lowercase ,       )
         a__       :						Tuple							=  model(__lowercase ,       attention_mask=__lowercase ,       token_type_ids=__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 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Union[str, Any]:
         """simple docstring"""




         a__       :						List[Any]							=  NezhaForMaskedLM(config=__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						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 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Optional[Any]:
         """simple docstring"""




         a__       :						List[str]							=  NezhaForNextSentencePrediction(config=__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						Dict							=  model(
             __lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       labels=__lowercase ,       )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, 2)  )

      def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Dict:
         """simple docstring"""




         a__       :						int							=  NezhaForPreTraining(config=__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						str							=  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 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Dict:
         """simple docstring"""




         a__       :						Optional[Any]							=  NezhaForQuestionAnswering(config=__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						Dict							=  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 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Any:
         """simple docstring"""




         a__       :						Optional[int]							=  self.num_labels
         a__       :						str							=  NezhaForSequenceClassification(__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						List[str]							=  model(__lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       labels=__lowercase  )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_labels)  )

      def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Any:
         """simple docstring"""




         a__       :						Union[str, Any]							=  self.num_labels
         a__       :						Tuple							=  NezhaForTokenClassification(config=__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						Optional[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 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase  ) ->				Optional[int]:
         """simple docstring"""




         a__       :						Union[str, Any]							=  self.num_choices
         a__       :						Dict							=  NezhaForMultipleChoice(config=__lowercase  )
         model.to(__lowercase  )
         model.eval()
         a__       :						Optional[int]							=  input_ids.unsqueeze(1  ).expand(-1 ,       self.num_choices ,       -1  ).contiguous()
         a__       :						Dict							=  token_type_ids.unsqueeze(1  ).expand(-1 ,       self.num_choices ,       -1  ).contiguous()
         a__       :						Optional[Any]							=  input_mask.unsqueeze(1  ).expand(-1 ,       self.num_choices ,       -1  ).contiguous()
         a__       :						Any							=  model(
             __lowercase ,       attention_mask=__lowercase ,       token_type_ids=__lowercase ,       labels=__lowercase ,       )
         self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.num_choices)  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Tuple:
         """simple docstring"""




         a__       :						str							=  self.prepare_config_and_inputs()
         (
             a__						
         )       :						Tuple							=  config_and_inputs
         a__       :						Optional[int]							=  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
         return config, inputs_dict




@require_torch
class        snake_case__    (A__ ,     A__ ,     A__ ,     unittest.TestCase		):







      """simple docstring"""




      __lowerCAmelCase					:Dict      =					(
          (
              NezhaModel,
              NezhaForMaskedLM,
              NezhaForMultipleChoice,
              NezhaForNextSentencePrediction,
              NezhaForPreTraining,
              NezhaForQuestionAnswering,
              NezhaForSequenceClassification,
              NezhaForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      __lowerCAmelCase					:int      =					(
          {
              "feature-extraction": NezhaModel,
              "fill-mask": NezhaForMaskedLM,
              "question-answering": NezhaForQuestionAnswering,
              "text-classification": NezhaForSequenceClassification,
              "token-classification": NezhaForTokenClassification,
              "zero-shot": NezhaForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __lowerCAmelCase					:Tuple      =					True

      def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase ,       __lowercase=False  ) ->				Optional[int]:
         """simple docstring"""




         a__       :						Dict							=  super()._prepare_for_class(__lowercase ,       __lowercase ,       return_labels=__lowercase  )

         if return_labels:
            if model_class in get_values(__lowercase  ):
               a__       :						str							=  torch.zeros(
                   (self.model_tester.batch_size, self.model_tester.seq_length) ,       dtype=torch.long ,       device=__lowercase  )
               a__       :						Dict							=  torch.zeros(
                   self.model_tester.batch_size ,       dtype=torch.long ,       device=__lowercase  )
         return inputs_dict

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Union[str, Any]:
         """simple docstring"""




         a__       :						Optional[Any]							=  NezhaModelTester(self  )
         a__       :						Tuple							=  ConfigTester(self ,       config_class=__lowercase ,       hidden_size=3_7  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Any:
         """simple docstring"""




         self.config_tester.run_common_tests()

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Dict:
         """simple docstring"""




         a__       :						Dict							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				int:
         """simple docstring"""




         a__       :						Dict							=  self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[int]:
         """simple docstring"""




         (
             a__						
         )       :						List[Any]							=  self.model_tester.prepare_config_and_inputs_for_decoder()

         a__       :						List[Any]							=  None

         self.model_tester.create_and_check_model_as_decoder(
             __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       __lowercase ,       )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				int:
         """simple docstring"""




         a__       :						List[str]							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Dict:
         """simple docstring"""




         a__       :						Optional[Any]							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				int:
         """simple docstring"""




         a__       :						Any							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_next_sequence_prediction(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Dict:
         """simple docstring"""




         a__       :						Tuple							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[str]:
         """simple docstring"""




         a__       :						str							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				str:
         """simple docstring"""




         a__       :						Optional[Any]							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__lowercase  )

      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[int]:
         """simple docstring"""




         a__       :						Optional[int]							=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__lowercase  )

      @slow
      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Optional[Any]:
         """simple docstring"""




         for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            a__       :						Tuple							=  NezhaModel.from_pretrained(__lowercase  )
            self.assertIsNotNone(__lowercase  )

      @slow
      @require_torch_gpu
      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				Dict:
         """simple docstring"""




         a__       :						Optional[int]							=  self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
            # NezhaForMultipleChoice behaves incorrectly in JIT environments.
            if model_class == NezhaForMultipleChoice:
               return

            a__       :						List[Any]							=  True
            a__       :						Tuple							=  model_class(config=__lowercase  )

            a__       :						Optional[Any]							=  self._prepare_for_class(__lowercase ,       __lowercase  )
            a__       :						str							=  torch.jit.trace(
                __lowercase ,       (inputs_dict["""input_ids"""].to("""cpu"""  ), inputs_dict["""attention_mask"""].to("""cpu"""  ))  )

            with tempfile.TemporaryDirectory() as tmp:
               torch.jit.save(__lowercase ,       os.path.join(__lowercase ,       """bert.pt"""  )  )
               a__       :						List[Any]							=  torch.jit.load(os.path.join(__lowercase ,       """bert.pt"""  ) ,       map_location=__lowercase  )
               loaded(inputs_dict["""input_ids"""].to(__lowercase  ) ,       inputs_dict["""attention_mask"""].to(__lowercase  )  )




@require_torch
class        snake_case__    (unittest.TestCase		):







      """simple docstring"""




      @slow
      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				List[str]:
         """simple docstring"""




         a__       :						List[str]							=  NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""  )
         a__       :						List[Any]							=  torch.tensor([[0, 1, 2, 3, 4, 5]]  )
         a__       :						Dict							=  torch.tensor([[0, 1, 1, 1, 1, 1]]  )
         with torch.no_grad():
            a__       :						int							=  model(__lowercase ,       attention_mask=__lowercase  )[0]
         a__       :						List[Any]							=  torch.Size((1, 6, 7_6_8)  )
         self.assertEqual(output.shape ,       __lowercase  )
         a__       :						Dict							=  torch.tensor([[[0.0_6_8_5, 0.2_4_4_1, 0.1_1_0_2], [0.0_6_0_0, 0.1_9_0_6, 0.1_3_4_9], [0.0_2_2_1, 0.0_8_1_9, 0.0_5_8_6]]]  )

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

      @slow
      def 							SCREAMING_SNAKE_CASE__(       self  ) ->				str:
         """simple docstring"""




         a__       :						List[Any]							=  NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""  )
         a__       :						List[Any]							=  torch.tensor([[0, 1, 2, 3, 4, 5]]  )
         a__       :						Any							=  torch.tensor([[1, 1, 1, 1, 1, 1]]  )
         with torch.no_grad():
            a__       :						List[str]							=  model(__lowercase ,       attention_mask=__lowercase  )[0]
         a__       :						Any							=  torch.Size((1, 6, 2_1_1_2_8)  )
         self.assertEqual(output.shape ,       __lowercase  )
         a__       :						Dict							=  torch.tensor(
             [[-2.7_9_3_9, -1.7_9_0_2, -2.2_1_8_9], [-2.8_5_8_5, -1.8_9_0_8, -2.3_7_2_3], [-2.6_4_9_9, -1.7_7_5_0, -2.2_5_5_8]]  )

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









def      lowerCAmelCase_					(  _lowercase    :       int)							->	int:







			"""simple docstring"""

			if not isinstance(_lowercase						,       _lowercase):
						raise TypeError("""only integers accepted as input""")
			else:
						a__       :						Any							=  str(abs(_lowercase))
						a__       :						str							=  [list(_lowercase) for char in range(len(_lowercase))]
						for index in range(len(_lowercase)):
									num_transpositions[index].pop(_lowercase)
						return max(
						    int("""""".join(list(_lowercase))) for transposition in num_transpositions)


if __name__ == "__main__":
		__import__("doctest").testmod()




