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  SCREAMING_SNAKE_CASE_ :
  '''simple docstring'''


  def __init__(				self    :					Any      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					int=13      ,	SCREAMING_SNAKE_CASE__    :					str=7      ,	SCREAMING_SNAKE_CASE__    :					List[str]=True      ,	SCREAMING_SNAKE_CASE__    :					Any=True      ,	SCREAMING_SNAKE_CASE__    :					int=True      ,	SCREAMING_SNAKE_CASE__    :					List[str]=True      ,	SCREAMING_SNAKE_CASE__    :					Dict=99      ,	SCREAMING_SNAKE_CASE__    :					Tuple=32      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]=5      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]=4      ,	SCREAMING_SNAKE_CASE__    :					int=37      ,	SCREAMING_SNAKE_CASE__    :					str="gelu"      ,	SCREAMING_SNAKE_CASE__    :					int=0.1      ,	SCREAMING_SNAKE_CASE__    :					Any=0.1      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]=1_28      ,	SCREAMING_SNAKE_CASE__    :					Any=32      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]=16      ,	SCREAMING_SNAKE_CASE__    :					Dict=2      ,	SCREAMING_SNAKE_CASE__    :					List[str]=0.0_2      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]=3      ,	SCREAMING_SNAKE_CASE__    :					str=4      ,	SCREAMING_SNAKE_CASE__    :					Dict=None      ,	)    ->				List[Any]:
         A						:     Union[str, Any]  =parent
         A						:     List[Any]  =batch_size
         A						:     Tuple  =seq_length
         A						:     List[str]  =is_training
         A						:     Optional[Any]  =use_input_mask
         A						:     Dict  =use_token_type_ids
         A						:     Any  =use_labels
         A						:     Tuple  =vocab_size
         A						:     str  =hidden_size
         A						:     List[str]  =num_hidden_layers
         A						:     Union[str, Any]  =num_attention_heads
         A						:     int  =intermediate_size
         A						:     str  =hidden_act
         A						:     List[str]  =hidden_dropout_prob
         A						:     int  =attention_probs_dropout_prob
         A						:     str  =max_position_embeddings
         A						:     Any  =type_vocab_size
         A						:     str  =type_sequence_label_size
         A						:     List[str]  =initializer_range
         A						:     Dict  =num_labels
         A						:     Dict  =num_choices
         A						:     Dict  =scope







  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[str]     )    ->				int:
         A						:     int  =ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size     )

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

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

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

         A						:     str  =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]     )    ->				List[Any]:
         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=_lowerCamelCase      ,	initializer_range=self.initializer_range      ,	)







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]     )    ->				Union[str, Any]:
         (
             (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  	
         )						:     Optional[Any]  =self.prepare_config_and_inputs()

         A						:     Optional[int]  =True
         A						:     Any  =floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
         A						:     Optional[int]  =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    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					int      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					str      ,	SCREAMING_SNAKE_CASE__    :					int      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					List[str]     )    ->				Optional[int]:
         A						:     Dict  =NezhaModel(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     Any  =model(_lowerCamelCase      ,	attention_mask=_lowerCamelCase      ,	token_type_ids=_lowerCamelCase     )
         A						:     str  =model(_lowerCamelCase      ,	token_type_ids=_lowerCamelCase     )
         A						:     Tuple  =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 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					Any      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					Any      ,	SCREAMING_SNAKE_CASE__    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					Any      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					str      ,	SCREAMING_SNAKE_CASE__    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	)    ->				Optional[Any]:
         A						:     int  =True
         A						:     str  =NezhaModel(_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     Optional[Any]  =model(
             _lowerCamelCase      ,	attention_mask=_lowerCamelCase      ,	token_type_ids=_lowerCamelCase      ,	encoder_hidden_states=_lowerCamelCase      ,	encoder_attention_mask=_lowerCamelCase      ,	)
         A						:     List[Any]  =model(
             _lowerCamelCase      ,	attention_mask=_lowerCamelCase      ,	token_type_ids=_lowerCamelCase      ,	encoder_hidden_states=_lowerCamelCase      ,	)
         A						:     Dict  =model(_lowerCamelCase      ,	attention_mask=_lowerCamelCase      ,	token_type_ids=_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 		SCREAMING_SNAKE_CASE_		(				self    :					Dict      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					str      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					List[str]     )    ->				List[str]:
         A						:     Tuple  =NezhaForMaskedLM(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     int  =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 		SCREAMING_SNAKE_CASE_		(				self    :					Any      ,	SCREAMING_SNAKE_CASE__    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					str      ,	SCREAMING_SNAKE_CASE__    :					List[Any]     )    ->				Optional[int]:
         A						:     Optional[int]  =NezhaForNextSentencePrediction(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     List[Any]  =model(
             _lowerCamelCase      ,	attention_mask=_lowerCamelCase      ,	token_type_ids=_lowerCamelCase      ,	labels=_lowerCamelCase      ,	)
         self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, 2)     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Any      ,	SCREAMING_SNAKE_CASE__    :					str      ,	SCREAMING_SNAKE_CASE__    :					Any      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					Dict      ,	SCREAMING_SNAKE_CASE__    :					Any     )    ->				List[str]:
         A						:     Optional[int]  =NezhaForPreTraining(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     Optional[Any]  =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 		SCREAMING_SNAKE_CASE_		(				self    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					int      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					int      ,	SCREAMING_SNAKE_CASE__    :					Tuple     )    ->				Union[str, Any]:
         A						:     List[str]  =NezhaForQuestionAnswering(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     Dict  =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 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					Dict      ,	SCREAMING_SNAKE_CASE__    :					int      ,	SCREAMING_SNAKE_CASE__    :					str      ,	SCREAMING_SNAKE_CASE__    :					int      ,	SCREAMING_SNAKE_CASE__    :					Dict      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					Tuple     )    ->				int:
         A						:     Union[str, Any]  =self.num_labels
         A						:     Dict  =NezhaForSequenceClassification(_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     int  =model(_lowerCamelCase      ,	attention_mask=_lowerCamelCase      ,	token_type_ids=_lowerCamelCase      ,	labels=_lowerCamelCase     )
         self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.num_labels)     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					Any      ,	SCREAMING_SNAKE_CASE__    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					Any     )    ->				Optional[int]:
         A						:     Tuple  =self.num_labels
         A						:     str  =NezhaForTokenClassification(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     Tuple  =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 		SCREAMING_SNAKE_CASE_		(				self    :					Any      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					Dict      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					Tuple      ,	SCREAMING_SNAKE_CASE__    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					List[str]     )    ->				Tuple:
         A						:     List[str]  =self.num_choices
         A						:     Union[str, Any]  =NezhaForMultipleChoice(config=_lowerCamelCase     )
         model.to(_lowerCamelCase     )
         model.eval()
         A						:     Optional[int]  =input_ids.unsqueeze(1     ).expand(-1      ,	self.num_choices      ,	-1     ).contiguous()
         A						:     Union[str, Any]  =token_type_ids.unsqueeze(1     ).expand(-1      ,	self.num_choices      ,	-1     ).contiguous()
         A						:     Tuple  =input_mask.unsqueeze(1     ).expand(-1      ,	self.num_choices      ,	-1     ).contiguous()
         A						:     Tuple  =model(
             _lowerCamelCase      ,	attention_mask=_lowerCamelCase      ,	token_type_ids=_lowerCamelCase      ,	labels=_lowerCamelCase      ,	)
         self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.num_choices)     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					int     )    ->				int:
         A						:     List[str]  =self.prepare_config_and_inputs()
         (
             (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  	
         )						:     Any  =config_and_inputs
         A						:     int  ={'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   ,						__UpperCAmelCase   ,						unittest.TestCase			):
  '''simple docstring'''


  lowercase		:     str            =				(
      (
          NezhaModel,
          NezhaForMaskedLM,
          NezhaForMultipleChoice,
          NezhaForNextSentencePrediction,
          NezhaForPreTraining,
          NezhaForQuestionAnswering,
          NezhaForSequenceClassification,
          NezhaForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  lowercase		:     List[str]            =				(
      {
          "feature-extraction": NezhaModel,
          "fill-mask": NezhaForMaskedLM,
          "question-answering": NezhaForQuestionAnswering,
          "text-classification": NezhaForSequenceClassification,
          "token-classification": NezhaForTokenClassification,
          "zero-shot": NezhaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  lowercase		:     List[Any]            =				True







  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]=False     )    ->				Optional[Any]:
         A						:     Tuple  =super()._prepare_for_class(_lowerCamelCase      ,	_lowerCamelCase      ,	return_labels=_lowerCamelCase     )

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







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Dict     )    ->				Any:
         A						:     Any  =NezhaModelTester(self     )
         A						:     int  =ConfigTester(self      ,	config_class=_lowerCamelCase      ,	hidden_size=37     )







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







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[Any]     )    ->				int:
         A						:     List[str]  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]     )    ->				Optional[Any]:
         A						:     Tuple  =self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Any     )    ->				List[str]:
         # This regression test was failing with PyTorch < 1.3
         (
             (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  (
             A	
         )  ,  	
         )						:     str  =self.model_tester.prepare_config_and_inputs_for_decoder()

         A						:     Dict  =None

         self.model_tester.create_and_check_model_as_decoder(
             _lowerCamelCase      ,	_lowerCamelCase      ,	_lowerCamelCase      ,	_lowerCamelCase      ,	_lowerCamelCase      ,	_lowerCamelCase      ,	_lowerCamelCase      ,	_lowerCamelCase      ,	_lowerCamelCase      ,	)







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Dict     )    ->				Any:
         A						:     Dict  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]     )    ->				str:
         A						:     str  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					str     )    ->				str:
         A						:     Optional[int]  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_next_sequence_prediction(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					str     )    ->				Optional[int]:
         A						:     Any  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Dict     )    ->				int:
         A						:     int  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Tuple     )    ->				Optional[Any]:
         A						:     List[str]  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Tuple     )    ->				Tuple:
         A						:     Union[str, Any]  =self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase     )







  @slow
  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[str]     )    ->				Dict:
         for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                A						:     Optional[int]  =NezhaModel.from_pretrained(_lowerCamelCase     )
                self.assertIsNotNone(_lowerCamelCase     )







  @slow
  @require_torch_gpu
  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[Any]     )    ->				Tuple:
         A  ,  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						:     str  =True
                A						:     Any  =model_class(config=_lowerCamelCase     )

                A						:     Union[str, Any]  =self._prepare_for_class(_lowerCamelCase      ,	_lowerCamelCase     )
                A						:     int  =torch.jit.trace(
                    _lowerCamelCase      ,	(inputs_dict['input_ids'].to('cpu'     ), inputs_dict['attention_mask'].to('cpu'     ))     )

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







@require_torch
class  SCREAMING_SNAKE_CASE_ (					unittest.TestCase			):
  '''simple docstring'''


  @slow
  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[str]     )    ->				List[Any]:
         A						:     Dict  =NezhaModel.from_pretrained('sijunhe/nezha-cn-base'     )
         A						:     Optional[Any]  =torch.tensor([[0, 1, 2, 3, 4, 5]]     )
         A						:     List[str]  =torch.tensor([[0, 1, 1, 1, 1, 1]]     )
         with torch.no_grad():
                A						:     List[str]  =model(_lowerCamelCase      ,	attention_mask=_lowerCamelCase     )[0]
         A						:     Tuple  =torch.Size((1, 6, 7_68)     )
         self.assertEqual(output.shape      ,	_lowerCamelCase     )
         A						:     Optional[int]  =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]      ,	_lowerCamelCase      ,	atol=1e-4     )     )







  @slow
  def 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]     )    ->				Dict:
         A						:     Optional[Any]  =NezhaForMaskedLM.from_pretrained('sijunhe/nezha-cn-base'     )
         A						:     Dict  =torch.tensor([[0, 1, 2, 3, 4, 5]]     )
         A						:     Optional[int]  =torch.tensor([[1, 1, 1, 1, 1, 1]]     )
         with torch.no_grad():
                A						:     Union[str, Any]  =model(_lowerCamelCase      ,	attention_mask=_lowerCamelCase     )[0]
         A						:     Optional[int]  =torch.Size((1, 6, 2_11_28)     )
         self.assertEqual(output.shape      ,	_lowerCamelCase     )
         A						:     int  =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]      ,	_lowerCamelCase      ,	atol=1e-4     )     )







from typing import List

from .keymap import KEYMAP, get_character



def A__  (			lowercase:     str					)				->		List[str]:

       def decorator(lowercase:     int					):
              A						:     Tuple         =getattr(lowercase,		'handle_key',		[]					)
              handle += [key]
              setattr(lowercase,		'handle_key',		lowercase					)
              return func

       return decorator



def A__  (			*lowercase:     List[str]					)				->		Dict:

       def decorator(lowercase:     Union[str, Any]					):
              A						:     Optional[int]         =getattr(lowercase,		'handle_key',		[]					)
              handle += keys
              setattr(lowercase,		'handle_key',		lowercase					)
              return func

       return decorator



class  SCREAMING_SNAKE_CASE_ (					lowerCAmelCase_			):
  '''simple docstring'''


  def __new__(				cls    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					Optional[Any]      ,	SCREAMING_SNAKE_CASE__    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					List[str]     )    ->				Any:
         A						:     Dict         =super().__new__(cls      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__     )
         if not hasattr(SCREAMING_SNAKE_CASE__      ,	'key_handler'     ):
                setattr(SCREAMING_SNAKE_CASE__      ,	'key_handler'      ,	{}     )
         setattr(SCREAMING_SNAKE_CASE__      ,	'handle_input'      ,	KeyHandler.handle_input     )

         for value in attrs.values():
                A						:     Optional[Any]         =getattr(SCREAMING_SNAKE_CASE__      ,	'handle_key'      ,	[]     )
                for key in handled_keys:
                       A						:     str         =value
         return new_cls







  @staticmethod
  def 		SCREAMING_SNAKE_CASE_		(				cls    :					str     )    ->				Any:
         A						:     str         =get_character()
         if char != KEYMAP["undefined"]:
                A						:     List[str]         =ord(SCREAMING_SNAKE_CASE__     )
         A						:     List[str]         =cls.key_handler.get(SCREAMING_SNAKE_CASE__     )
         if handler:
                A						:     List[str]         =char
                return handler(cls     )
         else:
                return None







def A__  (			cls:     Optional[int]					)				->		str:
       return KeyHandler(cls.__name__,		cls.__bases__,		cls.__dict__.copy()					)






