


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


					def __init__(    self						:		Union[str, Any]  ,     SCREAMING_SNAKE_CASE						:		Optional[Any]  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]=1_3  ,     SCREAMING_SNAKE_CASE						:		str=7  ,     SCREAMING_SNAKE_CASE						:		Optional[int]=True  ,     SCREAMING_SNAKE_CASE						:		int=True  ,     SCREAMING_SNAKE_CASE						:		Tuple=True  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]=True  ,     SCREAMING_SNAKE_CASE						:		str=9_9  ,     SCREAMING_SNAKE_CASE						:		List[str]=3_2  ,     SCREAMING_SNAKE_CASE						:		List[Any]=5  ,     SCREAMING_SNAKE_CASE						:		int=4  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]=3_7  ,     SCREAMING_SNAKE_CASE						:		List[Any]="gelu"  ,     SCREAMING_SNAKE_CASE						:		Optional[int]=0.1  ,     SCREAMING_SNAKE_CASE						:		Optional[Any]=0.1  ,     SCREAMING_SNAKE_CASE						:		str=1_2_8  ,     SCREAMING_SNAKE_CASE						:		str=3_2  ,     SCREAMING_SNAKE_CASE						:		Optional[Any]=1_6  ,     SCREAMING_SNAKE_CASE						:		List[str]=2  ,     SCREAMING_SNAKE_CASE						:		Tuple=0.02  ,     SCREAMING_SNAKE_CASE						:		int=3  ,     SCREAMING_SNAKE_CASE						:		int=4  ,     SCREAMING_SNAKE_CASE						:		Tuple=None  ,     )      ->						List[Any]:




											'''simple docstring'''

											__snake_case				=		parent
											__snake_case				=		batch_size
											__snake_case				=		seq_length
											__snake_case				=		is_training
											__snake_case				=		use_input_mask
											__snake_case				=		use_token_type_ids
											__snake_case				=		use_labels
											__snake_case				=		vocab_size
											__snake_case				=		hidden_size
											__snake_case				=		num_hidden_layers
											__snake_case				=		num_attention_heads
											__snake_case				=		intermediate_size
											__snake_case				=		hidden_act
											__snake_case				=		hidden_dropout_prob
											__snake_case				=		attention_probs_dropout_prob
											__snake_case				=		max_position_embeddings
											__snake_case				=		type_vocab_size
											__snake_case				=		type_sequence_label_size
											__snake_case				=		initializer_range
											__snake_case				=		num_labels
											__snake_case				=		num_choices
											__snake_case				=		scope


					def 				SCREAMING_SNAKE_CASE_	(    self						:		int    )      ->						Any:




											'''simple docstring'''

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

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

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

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

											__snake_case				=		self.get_config()

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


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple    )      ->						List[str]:




											'''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=__lowerCAmelCase  ,     initializer_range=self.initializer_range  ,     )


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




											'''simple docstring'''

											(
											    (
											    __snake_case       
											)		,       (
											    __snake_case       
											)		,       (
											    __snake_case       
											)		,       (
											    __snake_case       
											)		,       (
											    __snake_case       
											)		,       (
											    __snake_case       
											)		,       (
											    __snake_case       
											)		,              
											)				=		self.prepare_config_and_inputs()

											__snake_case				=		True
											__snake_case				=		floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
											__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 				SCREAMING_SNAKE_CASE_	(    self						:		Union[str, Any]  ,     SCREAMING_SNAKE_CASE						:		str  ,     SCREAMING_SNAKE_CASE						:		Any  ,     SCREAMING_SNAKE_CASE						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		Any  ,     SCREAMING_SNAKE_CASE						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		Tuple    )      ->						Any:




											'''simple docstring'''

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




											'''simple docstring'''

											__snake_case				=		True
											__snake_case				=		NezhaModel(__lowerCAmelCase    )
											model.to(__lowerCAmelCase    )
											model.eval()
											__snake_case				=		model(
											    __lowerCAmelCase  ,     attention_mask=__lowerCAmelCase  ,     token_type_ids=__lowerCAmelCase  ,     encoder_hidden_states=__lowerCAmelCase  ,     encoder_attention_mask=__lowerCAmelCase  ,     )
											__snake_case				=		model(
											    __lowerCAmelCase  ,     attention_mask=__lowerCAmelCase  ,     token_type_ids=__lowerCAmelCase  ,     encoder_hidden_states=__lowerCAmelCase  ,     )
											__snake_case				=		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						:		str  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]  ,     SCREAMING_SNAKE_CASE						:		Any  ,     SCREAMING_SNAKE_CASE						:		List[Any]  ,     SCREAMING_SNAKE_CASE						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]  ,     SCREAMING_SNAKE_CASE						:		List[Any]  ,     SCREAMING_SNAKE_CASE						:		Optional[Any]    )      ->						Tuple:




											'''simple docstring'''

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




											'''simple docstring'''

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




											'''simple docstring'''

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




											'''simple docstring'''

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




											'''simple docstring'''

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




											'''simple docstring'''

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




											'''simple docstring'''

											__snake_case				=		self.num_choices
											__snake_case				=		NezhaForMultipleChoice(config=__lowerCAmelCase    )
											model.to(__lowerCAmelCase    )
											model.eval()
											__snake_case				=		input_ids.unsqueeze(1    ).expand(-1  ,     self.num_choices  ,     -1    ).contiguous()
											__snake_case				=		token_type_ids.unsqueeze(1    ).expand(-1  ,     self.num_choices  ,     -1    ).contiguous()
											__snake_case				=		input_mask.unsqueeze(1    ).expand(-1  ,     self.num_choices  ,     -1    ).contiguous()
											__snake_case				=		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						:		Union[str, Any]    )      ->						int:




											'''simple docstring'''

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






@require_torch
class      UpperCamelCase(	UpperCAmelCase__  ,				UpperCAmelCase__  ,				UpperCAmelCase__  ,				unittest.TestCase ):
					snake_case_		:					List[Any]    					=						(
					    (
					        NezhaModel,
					        NezhaForMaskedLM,
					        NezhaForMultipleChoice,
					        NezhaForNextSentencePrediction,
					        NezhaForPreTraining,
					        NezhaForQuestionAnswering,
					        NezhaForSequenceClassification,
					        NezhaForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					snake_case_		:					int    					=						(
					    {
					        '''feature-extraction''': NezhaModel,
					        '''fill-mask''': NezhaForMaskedLM,
					        '''question-answering''': NezhaForQuestionAnswering,
					        '''text-classification''': NezhaForSequenceClassification,
					        '''token-classification''': NezhaForTokenClassification,
					        '''zero-shot''': NezhaForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					snake_case_		:					List[str]    					=						True


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str  ,     SCREAMING_SNAKE_CASE						:		int  ,     SCREAMING_SNAKE_CASE						:		List[Any]  ,     SCREAMING_SNAKE_CASE						:		Dict=False    )      ->						List[Any]:




											'''simple docstring'''

											__snake_case				=		super()._prepare_for_class(__lowerCAmelCase  ,     __lowerCAmelCase  ,     return_labels=__lowerCAmelCase    )

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


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str    )      ->						int:




											'''simple docstring'''

											__snake_case				=		NezhaModelTester(self    )
											__snake_case				=		ConfigTester(self  ,     config_class=__lowerCAmelCase  ,     hidden_size=3_7    )


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




											'''simple docstring'''

											self.config_tester.run_common_tests()


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




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__lowerCAmelCase    )


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




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*__lowerCAmelCase    )


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




											'''simple docstring'''

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

											__snake_case				=		None

											self.model_tester.create_and_check_model_as_decoder(
											    __lowerCAmelCase  ,     __lowerCAmelCase  ,     __lowerCAmelCase  ,     __lowerCAmelCase  ,     __lowerCAmelCase  ,     __lowerCAmelCase  ,     __lowerCAmelCase  ,     __lowerCAmelCase  ,     __lowerCAmelCase  ,     )


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




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		int    )      ->						Dict:




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*__lowerCAmelCase    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]    )      ->						Tuple:




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_next_sequence_prediction(*__lowerCAmelCase    )


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




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_pretraining(*__lowerCAmelCase    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple    )      ->						int:




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase    )


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




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase    )


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




											'''simple docstring'''

											__snake_case				=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase    )


					@slow
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Any    )      ->						Dict:




											'''simple docstring'''

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






					@slow
					@require_torch_gpu
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple    )      ->						Tuple:




											'''simple docstring'''

											__snake_case		,       __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

																	__snake_case				=		True
																	__snake_case				=		model_class(config=__lowerCAmelCase    )

																	__snake_case				=		self._prepare_for_class(__lowerCAmelCase  ,     __lowerCAmelCase    )
																	__snake_case				=		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"    )    )
																							__snake_case				=		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      UpperCamelCase(	unittest.TestCase ):


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




											'''simple docstring'''

											__snake_case				=		NezhaModel.from_pretrained("sijunhe/nezha-cn-base"    )
											__snake_case				=		torch.tensor([[0, 1, 2, 3, 4, 5]]    )
											__snake_case				=		torch.tensor([[0, 1, 1, 1, 1, 1]]    )
											with torch.no_grad():
																	__snake_case				=		model(__lowerCAmelCase  ,     attention_mask=__lowerCAmelCase    )[0]
											__snake_case				=		torch.Size((1, 6, 7_6_8)    )
											self.assertEqual(output.shape  ,     __lowerCAmelCase    )
											__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]  ,     __lowerCAmelCase  ,     atol=1e-4    )    )






					@slow
					def 				SCREAMING_SNAKE_CASE_	(    self						:		str    )      ->						Dict:




											'''simple docstring'''

											__snake_case				=		NezhaForMaskedLM.from_pretrained("sijunhe/nezha-cn-base"    )
											__snake_case				=		torch.tensor([[0, 1, 2, 3, 4, 5]]    )
											__snake_case				=		torch.tensor([[1, 1, 1, 1, 1, 1]]    )
											with torch.no_grad():
																	__snake_case				=		model(__lowerCAmelCase  ,     attention_mask=__lowerCAmelCase    )[0]
											__snake_case				=		torch.Size((1, 6, 2_1_1_2_8)    )
											self.assertEqual(output.shape  ,     __lowerCAmelCase    )
											__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]  ,     __lowerCAmelCase  ,     atol=1e-4    )    )











'''simple docstring'''


from .configuration_bert_masked import MaskedBertConfig
from .modeling_bert_masked import (
    MaskedBertForMultipleChoice,
    MaskedBertForQuestionAnswering,
    MaskedBertForSequenceClassification,
    MaskedBertForTokenClassification,
    MaskedBertModel,
)
from .modules import *

