



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

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

			def      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		Optional[Any]							      =     ids_tensor([self.batch_size, self.seq_length],   self.vocab_size					)

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

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

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

						UpperCAmelCase_:		List[Any]							      =     self.get_config()

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

			def      __snake_case       (self					)       ->   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=SCREAMING_SNAKE_CASE_,   initializer_range=self.initializer_range,   )

			def      __snake_case       (self					)       ->   Union[str, Any]:
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		Union[str, Any]							      =     self.prepare_config_and_inputs()

						UpperCAmelCase_:		Optional[Any]							      =     True
						UpperCAmelCase_:		Optional[Any]							      =     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
						UpperCAmelCase_:		Dict							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						UpperCAmelCase_:		int							      =     NezhaModel(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		int							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		str							      =     model(SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[str]							      =     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      __snake_case       (self,   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_,   )       ->   Tuple:
						UpperCAmelCase_:		Dict							      =     True
						UpperCAmelCase_:		Tuple							      =     NezhaModel(SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Optional[int]							      =     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_,   )
						UpperCAmelCase_:		List[str]							      =     model(
						    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   encoder_hidden_states=SCREAMING_SNAKE_CASE_,   )
						UpperCAmelCase_:		List[str]							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						UpperCAmelCase_:		Tuple							      =     NezhaForMaskedLM(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Optional[int]							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   int:
						UpperCAmelCase_:		Optional[Any]							      =     NezhaForNextSentencePrediction(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Dict							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						UpperCAmelCase_:		Tuple							      =     NezhaForPreTraining(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Tuple							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Union[str, Any]:
						UpperCAmelCase_:		Optional[Any]							      =     NezhaForQuestionAnswering(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Any							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[Any]:
						UpperCAmelCase_:		List[str]							      =     self.num_labels
						UpperCAmelCase_:		str							      =     NezhaForSequenceClassification(SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Dict							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   int:
						UpperCAmelCase_:		int							      =     self.num_labels
						UpperCAmelCase_:		List[str]							      =     NezhaForTokenClassification(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		List[Any]							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						UpperCAmelCase_:		Tuple							      =     self.num_choices
						UpperCAmelCase_:		Union[str, Any]							      =     NezhaForMultipleChoice(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Dict							      =     input_ids.unsqueeze(1					).expand(-1,   self.num_choices,   -1					).contiguous()
						UpperCAmelCase_:		Any							      =     token_type_ids.unsqueeze(1					).expand(-1,   self.num_choices,   -1					).contiguous()
						UpperCAmelCase_:		int							      =     input_mask.unsqueeze(1					).expand(-1,   self.num_choices,   -1					).contiguous()
						UpperCAmelCase_:		int							      =     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      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		Tuple							      =     self.prepare_config_and_inputs()
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		Union[str, Any]							      =     config_and_inputs
						UpperCAmelCase_:		List[str]							      =     {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict







@require_torch
class 				_a				(   _lowerCAmelCase ,		_lowerCAmelCase ,		_lowerCAmelCase ,		unittest.TestCase					):
			A      		=	(
			    (
			        NezhaModel,
			        NezhaForMaskedLM,
			        NezhaForMultipleChoice,
			        NezhaForNextSentencePrediction,
			        NezhaForPreTraining,
			        NezhaForQuestionAnswering,
			        NezhaForSequenceClassification,
			        NezhaForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			A      		=	(
			    {
			        '''feature-extraction''': NezhaModel,
			        '''fill-mask''': NezhaForMaskedLM,
			        '''question-answering''': NezhaForQuestionAnswering,
			        '''text-classification''': NezhaForSequenceClassification,
			        '''token-classification''': NezhaForTokenClassification,
			        '''zero-shot''': NezhaForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			A      		=	True

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=False					)       ->   Optional[int]:
						UpperCAmelCase_:		Dict							      =     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_					):
												UpperCAmelCase_:		Optional[int]							      =     torch.zeros(
												    (self.model_tester.batch_size, self.model_tester.seq_length),   dtype=torch.long,   device=SCREAMING_SNAKE_CASE_					)
												UpperCAmelCase_:		Any							      =     torch.zeros(
												    self.model_tester.batch_size,   dtype=torch.long,   device=SCREAMING_SNAKE_CASE_					)
						return inputs_dict

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		List[str]							      =     NezhaModelTester(self					)
						UpperCAmelCase_:		List[str]							      =     ConfigTester(self,   config_class=SCREAMING_SNAKE_CASE_,   hidden_size=37					)

			def      __snake_case       (self					)       ->   str:
						self.config_tester.run_common_tests()

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		Dict							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		List[Any]							      =     self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Tuple:
						# This regression test was failing with PyTorch < 1.3
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		Any							      =     self.model_tester.prepare_config_and_inputs_for_decoder()

						UpperCAmelCase_:		str							      =     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      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		List[Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		Union[str, Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		Dict							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_next_sequence_prediction(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		Optional[int]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		Any							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		Union[str, Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		Dict							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_					)

			@slow
			def      __snake_case       (self					)       ->   Tuple:
						for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCAmelCase_:		str							      =     NezhaModel.from_pretrained(SCREAMING_SNAKE_CASE_					)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_					)

			@slow
			@require_torch_gpu
			def      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_      ,							UpperCAmelCase_:		List[str]							      =     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

									UpperCAmelCase_:		Union[str, Any]							      =     True
									UpperCAmelCase_:		Any							      =     model_class(config=SCREAMING_SNAKE_CASE_					)

									UpperCAmelCase_:		Optional[int]							      =     self._prepare_for_class(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)
									UpperCAmelCase_:		List[Any]							      =     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"""					)					)
												UpperCAmelCase_:		Optional[Any]							      =     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 				_a				(   unittest.TestCase					):

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

			@slow
			def      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		Tuple							      =     NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""					)
						UpperCAmelCase_:		Tuple							      =     torch.tensor([[0, 1, 2, 3, 4, 5]]					)
						UpperCAmelCase_:		Optional[Any]							      =     torch.tensor([[1, 1, 1, 1, 1, 1]]					)
						with torch.no_grad():
									UpperCAmelCase_:		str							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_					)[0]
						UpperCAmelCase_:		str							      =     torch.Size((1, 6, 21128)					)
						self.assertEqual(output.shape,   SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     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],   SCREAMING_SNAKE_CASE_,   atol=1E-4					)					)











from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


a      :  Optional[Any]			     =			{
    'configuration_mvp': ['MVP_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MvpConfig', 'MvpOnnxConfig'],
    'tokenization_mvp': ['MvpTokenizer'],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	a      :  str			     =			['MvpTokenizerFast']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	a      :  Optional[Any]			     =			[
	    'MVP_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'MvpForCausalLM',
	    'MvpForConditionalGeneration',
	    'MvpForQuestionAnswering',
	    'MvpForSequenceClassification',
	    'MvpModel',
	    'MvpPreTrainedModel',
	]

if TYPE_CHECKING:
	from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
	from .tokenization_mvp import MvpTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_mvp_fast import MvpTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_mvp import (
		    MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MvpForCausalLM,
		    MvpForConditionalGeneration,
		    MvpForQuestionAnswering,
		    MvpForSequenceClassification,
		    MvpModel,
		    MvpPreTrainedModel,
		)

else:
	import sys

	a      :  Optional[Any]			     =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






