

'''simple docstring'''



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



 """simple docstring"""



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


      """simple docstring"""




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


      """simple docstring"""




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

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

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

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

      __SCREAMING_SNAKE_CASE     =      self.get_config()

      return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
 def       UpperCAmelCase__		(		self					:       int			)  -> 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=__SCREAMING_SNAKE_CASE					,							initializer_range=self.initializer_range					,							)
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> str:


      """simple docstring"""




      (
          (
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			
      )     =      self.prepare_config_and_inputs()

      __SCREAMING_SNAKE_CASE     =      True
      __SCREAMING_SNAKE_CASE     =      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
      __SCREAMING_SNAKE_CASE     =      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       UpperCAmelCase__		(		self					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       Any					,							__SCREAMING_SNAKE_CASE					:       Dict					,							__SCREAMING_SNAKE_CASE					:       Any					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       Dict					,							__SCREAMING_SNAKE_CASE					:       int			)  -> str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaModel(config=__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE					,							attention_mask=__SCREAMING_SNAKE_CASE					,							token_type_ids=__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE					,							token_type_ids=__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      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					:       Dict					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]					,							__SCREAMING_SNAKE_CASE					:       Dict					,							__SCREAMING_SNAKE_CASE					:       Tuple					,							__SCREAMING_SNAKE_CASE					:       Dict					,							__SCREAMING_SNAKE_CASE					:       str					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       Dict					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       Any					,							)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      True
      __SCREAMING_SNAKE_CASE     =      NezhaModel(__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      model(
          __SCREAMING_SNAKE_CASE					,							attention_mask=__SCREAMING_SNAKE_CASE					,							token_type_ids=__SCREAMING_SNAKE_CASE					,							encoder_hidden_states=__SCREAMING_SNAKE_CASE					,							encoder_attention_mask=__SCREAMING_SNAKE_CASE					,							)
      __SCREAMING_SNAKE_CASE     =      model(
          __SCREAMING_SNAKE_CASE					,							attention_mask=__SCREAMING_SNAKE_CASE					,							token_type_ids=__SCREAMING_SNAKE_CASE					,							encoder_hidden_states=__SCREAMING_SNAKE_CASE					,							)
      __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE					,							attention_mask=__SCREAMING_SNAKE_CASE					,							token_type_ids=__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					:       Dict					,							__SCREAMING_SNAKE_CASE					:       str					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       str			)  -> Optional[int]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaForMaskedLM(config=__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      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					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       str					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]					,							__SCREAMING_SNAKE_CASE					:       Any					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       str					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       List[str]			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaForNextSentencePrediction(config=__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      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					:       Tuple					,							__SCREAMING_SNAKE_CASE					:       Any					,							__SCREAMING_SNAKE_CASE					:       str					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       str					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       List[str]			)  -> Dict:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaForPreTraining(config=__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      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					:       List[str]					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       Tuple					,							__SCREAMING_SNAKE_CASE					:       List[str]					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]					,							__SCREAMING_SNAKE_CASE					:       Any					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]			)  -> Optional[int]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaForQuestionAnswering(config=__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      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					:       int					,							__SCREAMING_SNAKE_CASE					:       List[str]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]					,							__SCREAMING_SNAKE_CASE					:       Any					,							__SCREAMING_SNAKE_CASE					:       List[str]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]					,							__SCREAMING_SNAKE_CASE					:       Optional[int]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]			)  -> Any:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.num_labels
      __SCREAMING_SNAKE_CASE     =      NezhaForSequenceClassification(__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      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					:       int					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       Optional[int]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       int			)  -> Any:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.num_labels
      __SCREAMING_SNAKE_CASE     =      NezhaForTokenClassification(config=__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      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					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       str					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       int					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       Dict					,							__SCREAMING_SNAKE_CASE					:       List[str]					,							__SCREAMING_SNAKE_CASE					:       int			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.num_choices
      __SCREAMING_SNAKE_CASE     =      NezhaForMultipleChoice(config=__SCREAMING_SNAKE_CASE			)
      model.to(__SCREAMING_SNAKE_CASE			)
      model.eval()
      __SCREAMING_SNAKE_CASE     =      input_ids.unsqueeze(1			).expand(-1					,							self.num_choices					,							-1			).contiguous()
      __SCREAMING_SNAKE_CASE     =      token_type_ids.unsqueeze(1			).expand(-1					,							self.num_choices					,							-1			).contiguous()
      __SCREAMING_SNAKE_CASE     =      input_mask.unsqueeze(1			).expand(-1					,							self.num_choices					,							-1			).contiguous()
      __SCREAMING_SNAKE_CASE     =      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					:       Optional[int]			)  -> Dict:


      """simple docstring"""




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







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



 """simple docstring"""



 lowerCAmelCase__  					=      (
     (
         NezhaModel,
         NezhaForMaskedLM,
         NezhaForMultipleChoice,
         NezhaForNextSentencePrediction,
         NezhaForPreTraining,
         NezhaForQuestionAnswering,
         NezhaForSequenceClassification,
         NezhaForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 lowerCAmelCase__  					=      (
     {
         "feature-extraction": NezhaModel,
         "fill-mask": NezhaForMaskedLM,
         "question-answering": NezhaForQuestionAnswering,
         "text-classification": NezhaForSequenceClassification,
         "token-classification": NezhaForTokenClassification,
         "zero-shot": NezhaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 lowerCAmelCase__  					=      True
 def       UpperCAmelCase__		(		self					:       Dict					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]					,							__SCREAMING_SNAKE_CASE					:       List[Any]					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]=False			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      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			):
                __SCREAMING_SNAKE_CASE     =      torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.seq_length)					,							dtype=torch.long					,							device=__SCREAMING_SNAKE_CASE			)
                __SCREAMING_SNAKE_CASE     =      torch.zeros(
                    self.model_tester.batch_size					,							dtype=torch.long					,							device=__SCREAMING_SNAKE_CASE			)
      return inputs_dict
 def       UpperCAmelCase__		(		self					:       Tuple			)  -> str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaModelTester(self			)
      __SCREAMING_SNAKE_CASE     =      ConfigTester(self					,							config_class=__SCREAMING_SNAKE_CASE					,							hidden_size=37			)
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> Optional[int]:


      """simple docstring"""




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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Optional[int]			)  -> Union[str, Any]:


      """simple docstring"""




      (
          (
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			(
          __SCREAMING_SNAKE_CASE
      )   ,			
      )     =      self.model_tester.prepare_config_and_inputs_for_decoder()

      __SCREAMING_SNAKE_CASE     =      None

      self.model_tester.create_and_check_model_as_decoder(
          __SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE					,							)
 def       UpperCAmelCase__		(		self					:       Dict			)  -> int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Any			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       int			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_next_sequence_prediction(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Dict			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> Optional[int]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       int			)  -> Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*__SCREAMING_SNAKE_CASE			)
 @slow
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Tuple:


      """simple docstring"""




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


 @slow
 @require_torch_gpu
 def       UpperCAmelCase__		(		self					:       Tuple			)  -> List[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      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

           __SCREAMING_SNAKE_CASE     =      True
           __SCREAMING_SNAKE_CASE     =      model_class(config=__SCREAMING_SNAKE_CASE			)

           __SCREAMING_SNAKE_CASE     =      self._prepare_for_class(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
           __SCREAMING_SNAKE_CASE     =      torch.jit.trace(
               __SCREAMING_SNAKE_CASE					,							(inputs_dict["""input_ids"""].to("""cpu"""			), inputs_dict["""attention_mask"""].to("""cpu"""			))			)

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







@require_torch
class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



 @slow
 def       UpperCAmelCase__		(		self					:       Union[str, Any]			)  -> Any:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""			)
      __SCREAMING_SNAKE_CASE     =      torch.tensor([[0, 1, 2, 3, 4, 5]]			)
      __SCREAMING_SNAKE_CASE     =      torch.tensor([[0, 1, 1, 1, 1, 1]]			)
      with torch.no_grad():
           __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE					,							attention_mask=__SCREAMING_SNAKE_CASE			)[0]
      __SCREAMING_SNAKE_CASE     =      torch.Size((1, 6, 768)			)
      self.assertEqual(output.shape					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]			)

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


 @slow
 def       UpperCAmelCase__		(		self					:       Tuple			)  -> Tuple:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""			)
      __SCREAMING_SNAKE_CASE     =      torch.tensor([[0, 1, 2, 3, 4, 5]]			)
      __SCREAMING_SNAKE_CASE     =      torch.tensor([[1, 1, 1, 1, 1, 1]]			)
      with torch.no_grad():
           __SCREAMING_SNAKE_CASE     =      model(__SCREAMING_SNAKE_CASE					,							attention_mask=__SCREAMING_SNAKE_CASE			)[0]
      __SCREAMING_SNAKE_CASE     =      torch.Size((1, 6, 21_128)			)
      self.assertEqual(output.shape					,							__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      torch.tensor(
          [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]			)

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









'''simple docstring'''



import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


UpperCAmelCase :       str			   =     logging.get_logger(__name__)







class 						lowerCAmelCase__   (      a					):



 """simple docstring"""



 lowerCAmelCase__  					=      "linear"
 lowerCAmelCase__  					=      "cosine"
 lowerCAmelCase__  					=      "cosine_with_restarts"
 lowerCAmelCase__  					=      "polynomial"
 lowerCAmelCase__  					=      "constant"
 lowerCAmelCase__  					=      "constant_with_warmup"
 lowerCAmelCase__  					=      "piecewise_constant"







def    a__       (  a__      ,    a__ = -1						):




     """simple docstring"""


     return LambdaLR(a__      ,    lambda a__						: 1      ,    last_epoch=a__						)







def    a__       (  a__      ,    a__      ,    a__ = -1						):




     """simple docstring"""


     def lr_lambda(a__						):
          if current_step < num_warmup_steps:
               return float(a__						) / float(max(1.0      ,    a__						)						)
          return 1.0

     return LambdaLR(a__      ,    a__      ,    last_epoch=a__						)







def    a__       (  a__      ,    a__      ,    a__ = -1						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      {}
     __SCREAMING_SNAKE_CASE     =      step_rules.split(""","""						)
     for rule_str in rule_list[:-1]:
          __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      rule_str.split(""":"""						)
          __SCREAMING_SNAKE_CASE     =      int(a__						)
          __SCREAMING_SNAKE_CASE     =      float(a__						)
          __SCREAMING_SNAKE_CASE     =      value
     __SCREAMING_SNAKE_CASE     =      float(rule_list[-1]						)

     def create_rules_function(a__      ,    a__						):
          def rule_func(a__						) -> float:
               __SCREAMING_SNAKE_CASE     =      sorted(rules_dict.keys()						)
               for i, sorted_step in enumerate(a__						):
                    if steps < sorted_step:
                         return rules_dict[sorted_steps[i]]
               return last_lr_multiple

          return rule_func

     __SCREAMING_SNAKE_CASE     =      create_rules_function(a__      ,    a__						)

     return LambdaLR(a__      ,    a__      ,    last_epoch=a__						)







def    a__       (  a__      ,    a__      ,    a__      ,    a__=-1						):




     """simple docstring"""


     def lr_lambda(a__						):
          if current_step < num_warmup_steps:
               return float(a__						) / float(max(1      ,    a__						)						)
          return max(
              0.0      ,    float(num_training_steps - current_step						) / float(max(1      ,    num_training_steps - num_warmup_steps						)						)						)

     return LambdaLR(a__      ,    a__      ,    a__						)







def    a__       (  a__      ,    a__      ,    a__      ,    a__ = 0.5      ,    a__ = -1						):




     """simple docstring"""


     def lr_lambda(a__						):
          if current_step < num_warmup_steps:
               return float(a__						) / float(max(1      ,    a__						)						)
          __SCREAMING_SNAKE_CASE     =      float(current_step - num_warmup_steps						) / float(max(1      ,    num_training_steps - num_warmup_steps						)						)
          return max(0.0      ,    0.5 * (1.0 + math.cos(math.pi * float(a__						) * 2.0 * progress						))						)

     return LambdaLR(a__      ,    a__      ,    a__						)







def    a__       (  a__      ,    a__      ,    a__      ,    a__ = 1      ,    a__ = -1						):




     """simple docstring"""


     def lr_lambda(a__						):
          if current_step < num_warmup_steps:
               return float(a__						) / float(max(1      ,    a__						)						)
          __SCREAMING_SNAKE_CASE     =      float(current_step - num_warmup_steps						) / float(max(1      ,    num_training_steps - num_warmup_steps						)						)
          if progress >= 1.0:
               return 0.0
          return max(0.0      ,    0.5 * (1.0 + math.cos(math.pi * ((float(a__						) * progress) % 1.0)						))						)

     return LambdaLR(a__      ,    a__      ,    a__						)







def    a__       (  a__      ,    a__      ,    a__      ,    a__=1E-7      ,    a__=1.0      ,    a__=-1						):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      optimizer.defaults["""lr"""]
     if not (lr_init > lr_end):
          raise ValueError(F'lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})'						)

     def lr_lambda(a__						):
          if current_step < num_warmup_steps:
               return float(a__						) / float(max(1      ,    a__						)						)
          elif current_step > num_training_steps:
               return lr_end / lr_init  # as LambdaLR multiplies by lr_init
          else:
               __SCREAMING_SNAKE_CASE     =      lr_init - lr_end
               __SCREAMING_SNAKE_CASE     =      num_training_steps - num_warmup_steps
               __SCREAMING_SNAKE_CASE     =      1 - (current_step - num_warmup_steps) / decay_steps
               __SCREAMING_SNAKE_CASE     =      lr_range * pct_remaining**power + lr_end
               return decay / lr_init  # as LambdaLR multiplies by lr_init

     return LambdaLR(a__      ,    a__      ,    a__						)


UpperCAmelCase :       Optional[Any]			   =     {
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}







def    a__       (  a__      ,    a__      ,    a__ = None      ,    a__ = None      ,    a__ = None      ,    a__ = 1      ,    a__ = 1.0      ,    a__ = -1      ,    ):




     """simple docstring"""


     __SCREAMING_SNAKE_CASE     =      SchedulerType(a__						)
     __SCREAMING_SNAKE_CASE     =      TYPE_TO_SCHEDULER_FUNCTION[name]
     if name == SchedulerType.CONSTANT:
          return schedule_func(a__      ,    last_epoch=a__						)

     if name == SchedulerType.PIECEWISE_CONSTANT:
          return schedule_func(a__      ,    step_rules=a__      ,    last_epoch=a__						)

     # All other schedulers require `num_warmup_steps`
     if num_warmup_steps is None:
          raise ValueError(F'{name} requires `num_warmup_steps`, please provide that argument.'						)

     if name == SchedulerType.CONSTANT_WITH_WARMUP:
          return schedule_func(a__      ,    num_warmup_steps=a__      ,    last_epoch=a__						)

     # All other schedulers require `num_training_steps`
     if num_training_steps is None:
          raise ValueError(F'{name} requires `num_training_steps`, please provide that argument.'						)

     if name == SchedulerType.COSINE_WITH_RESTARTS:
          return schedule_func(
              a__      ,    num_warmup_steps=a__      ,    num_training_steps=a__      ,    num_cycles=a__      ,    last_epoch=a__      ,    )

     if name == SchedulerType.POLYNOMIAL:
          return schedule_func(
              a__      ,    num_warmup_steps=a__      ,    num_training_steps=a__      ,    power=a__      ,    last_epoch=a__      ,    )

     return schedule_func(
         a__      ,    num_warmup_steps=a__      ,    num_training_steps=a__      ,    last_epoch=a__						)






