





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			, _lowerCAmelCase			, _lowerCAmelCase=13			, _lowerCAmelCase=7			, _lowerCAmelCase=True			, _lowerCAmelCase=True			, _lowerCAmelCase=True			, _lowerCAmelCase=True			, _lowerCAmelCase=99			, _lowerCAmelCase=32			, _lowerCAmelCase=5			, _lowerCAmelCase=4			, _lowerCAmelCase=37			, _lowerCAmelCase="gelu"			, _lowerCAmelCase=0.1			, _lowerCAmelCase=0.1			, _lowerCAmelCase=128			, _lowerCAmelCase=32			, _lowerCAmelCase=16			, _lowerCAmelCase=2			, _lowerCAmelCase=0.02			, _lowerCAmelCase=3			, _lowerCAmelCase=4			, _lowerCAmelCase=None			, ):

										"""simple docstring"""




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






				def 			snake_case_    (  self					):

										"""simple docstring"""




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

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

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

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

										__SCREAMING_SNAKE_CASE:				str           =   self.get_config()

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






				def 			snake_case_    (  self					):

										"""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 			snake_case_    (  self					):

										"""simple docstring"""




										(
										    (
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,
										):				List[Any]           =   self.prepare_config_and_inputs()

										__SCREAMING_SNAKE_CASE:				str           =   True
										__SCREAMING_SNAKE_CASE:				int           =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
										__SCREAMING_SNAKE_CASE:				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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




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

										"""simple docstring"""




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

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Tuple           =   NezhaForMaskedLM(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   NezhaForNextSentencePrediction(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				Any           =   model(
										    _lowerCAmelCase			, attention_mask=_lowerCAmelCase			, token_type_ids=_lowerCAmelCase			, labels=_lowerCAmelCase			, )
										self.parent.assertEqual(result.logits.shape			, (self.batch_size, 2)					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Any           =   NezhaForPreTraining(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




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

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				str           =   self.num_labels
										__SCREAMING_SNAKE_CASE:				List[Any]           =   NezhaForSequenceClassification(_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				Tuple           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase			, token_type_ids=_lowerCAmelCase			, labels=_lowerCAmelCase					)
										self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.num_labels)					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




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

										"""simple docstring"""




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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   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
										)       ,
										):				List[str]           =   config_and_inputs
										__SCREAMING_SNAKE_CASE:				str           =   {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
										return config, inputs_dict





@require_torch
class 				a			(      __lowercase   ,__lowercase   ,__lowercase   ,unittest.TestCase			):
				SCREAMING_SNAKE_CASE__     :   Optional[Any]		      =     (
				    (
				        NezhaModel,
				        NezhaForMaskedLM,
				        NezhaForMultipleChoice,
				        NezhaForNextSentencePrediction,
				        NezhaForPreTraining,
				        NezhaForQuestionAnswering,
				        NezhaForSequenceClassification,
				        NezhaForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				SCREAMING_SNAKE_CASE__     :   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 {}
				)
				SCREAMING_SNAKE_CASE__     :   Optional[int]		      =     True






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase=False					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   super()._prepare_for_class(_lowerCAmelCase			, _lowerCAmelCase			, return_labels=_lowerCAmelCase					)

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   NezhaModelTester(self					)
										__SCREAMING_SNAKE_CASE:				str           =   ConfigTester(self			, config_class=_lowerCAmelCase			, hidden_size=37					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										self.config_tester.run_common_tests()






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_model_as_decoder(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""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
										)       ,
										):				str           =   self.model_tester.prepare_config_and_inputs_for_decoder()

										__SCREAMING_SNAKE_CASE:				Dict           =   None

										self.model_tester.create_and_check_model_as_decoder(
										    _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, )






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_next_sequence_prediction(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_pretraining(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Dict           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				int           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Dict           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase					)






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




										for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   NezhaModel.from_pretrained(_lowerCAmelCase					)
																self.assertIsNotNone(_lowerCAmelCase					)






				@slow
				@require_torch_gpu
				def 			snake_case_    (  self					):

										"""simple docstring"""




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

																__SCREAMING_SNAKE_CASE:				Dict           =   True
																__SCREAMING_SNAKE_CASE:				Any           =   model_class(config=_lowerCAmelCase					)

																__SCREAMING_SNAKE_CASE:				List[str]           =   self._prepare_for_class(_lowerCAmelCase			, _lowerCAmelCase					)
																__SCREAMING_SNAKE_CASE:				Optional[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'''					)					)
																						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   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 				a			(      unittest.TestCase			):






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   NezhaModel.from_pretrained('''sijunhe/nezha-cn-base'''					)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   torch.tensor([[0, 1, 2, 3, 4, 5]]					)
										__SCREAMING_SNAKE_CASE:				int           =   torch.tensor([[0, 1, 1, 1, 1, 1]]					)
										with torch.no_grad():
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase					)[0]
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.Size((1, 6, 768)					)
										self.assertEqual(output.shape			, _lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				List[Any]           =   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 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   NezhaForMaskedLM.from_pretrained('''sijunhe/nezha-cn-base'''					)
										__SCREAMING_SNAKE_CASE:				List[Any]           =   torch.tensor([[0, 1, 2, 3, 4, 5]]					)
										__SCREAMING_SNAKE_CASE:				str           =   torch.tensor([[1, 1, 1, 1, 1, 1]]					)
										with torch.no_grad():
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase					)[0]
										__SCREAMING_SNAKE_CASE:				int           =   torch.Size((1, 6, 21128)					)
										self.assertEqual(output.shape			, _lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Dict           =   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					)					)










from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
							import PIL


lowerCAmelCase							:       str					 =	logging.get_logger(__name__)





class 				a			(      __lowercase			):
				SCREAMING_SNAKE_CASE__     :   Any		      =     ['''pixel_values''']






				def __init__(  self			, _lowerCAmelCase = True			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = PILImageResampling.BILINEAR			, _lowerCAmelCase = True			, _lowerCAmelCase = 1 / 255			, _lowerCAmelCase = True			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										super().__init__(**_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   size if size is not None else {'''shortest_edge''': 384}
										__SCREAMING_SNAKE_CASE:				List[str]           =   get_size_dict(_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   do_resize
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   size
										# Default value set here for backwards compatibility where the value in config is None
										__SCREAMING_SNAKE_CASE:				List[str]           =   crop_pct if crop_pct is not None else 224 / 256
										__SCREAMING_SNAKE_CASE:				List[Any]           =   resample
										__SCREAMING_SNAKE_CASE:				Any           =   do_rescale
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   rescale_factor
										__SCREAMING_SNAKE_CASE:				Any           =   do_normalize
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   image_std if image_std is not None else IMAGENET_STANDARD_STD






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase = PILImageResampling.BICUBIC			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   get_size_dict(_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)
										if "shortest_edge" not in size:
																raise ValueError(f"""Size dictionary must contain 'shortest_edge' key. Got {size.keys()}"""					)
										__SCREAMING_SNAKE_CASE:				str           =   size['''shortest_edge''']

										if shortest_edge < 384:
																# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
																__SCREAMING_SNAKE_CASE:				Dict           =   int(shortest_edge / crop_pct					)
																__SCREAMING_SNAKE_CASE:				Any           =   get_resize_output_image_size(_lowerCAmelCase			, size=_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   resize(image=_lowerCAmelCase			, size=_lowerCAmelCase			, resample=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)
																# then crop to (shortest_edge, shortest_edge)
																return center_crop(image=_lowerCAmelCase			, size=(shortest_edge, shortest_edge)			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)
										else:
																# warping (no cropping) when evaluated at 384 or larger
																return resize(
																    _lowerCAmelCase			, size=(shortest_edge, shortest_edge)			, resample=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										return rescale(_lowerCAmelCase			, scale=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										return normalize(_lowerCAmelCase			, mean=_lowerCAmelCase			, std=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = ChannelDimension.FIRST			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   do_resize if do_resize is not None else self.do_resize
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   crop_pct if crop_pct is not None else self.crop_pct
										__SCREAMING_SNAKE_CASE:				Tuple           =   resample if resample is not None else self.resample
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   do_rescale if do_rescale is not None else self.do_rescale
										__SCREAMING_SNAKE_CASE:				Any           =   rescale_factor if rescale_factor is not None else self.rescale_factor
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   do_normalize if do_normalize is not None else self.do_normalize
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   image_mean if image_mean is not None else self.image_mean
										__SCREAMING_SNAKE_CASE:				str           =   image_std if image_std is not None else self.image_std

										__SCREAMING_SNAKE_CASE:				List[Any]           =   size if size is not None else self.size
										__SCREAMING_SNAKE_CASE:				Dict           =   get_size_dict(_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				List[str]           =   make_list_of_images(_lowerCAmelCase					)

										if not valid_images(_lowerCAmelCase					):
																raise ValueError(
																    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
																    '''torch.Tensor, tf.Tensor or jax.ndarray.'''					)

										if do_resize and size is None or resample is None:
																raise ValueError('''Size and resample must be specified if do_resize is True.'''					)

										if do_resize and size["shortest_edge"] < 384 and crop_pct is None:
																raise ValueError('''crop_pct must be specified if size < 384.'''					)

										if do_rescale and rescale_factor is None:
																raise ValueError('''Rescale factor must be specified if do_rescale is True.'''					)

										if do_normalize and (image_mean is None or image_std is None):
																raise ValueError('''Image mean and std must be specified if do_normalize is True.'''					)

										# All transformations expect numpy arrays.
										__SCREAMING_SNAKE_CASE:				Tuple           =   [to_numpy_array(_lowerCAmelCase					) for image in images]

										if do_resize:
																__SCREAMING_SNAKE_CASE:				str           =   [self.resize(image=_lowerCAmelCase			, size=_lowerCAmelCase			, crop_pct=_lowerCAmelCase			, resample=_lowerCAmelCase					) for image in images]

										if do_rescale:
																__SCREAMING_SNAKE_CASE:				Any           =   [self.rescale(image=_lowerCAmelCase			, scale=_lowerCAmelCase					) for image in images]

										if do_normalize:
																__SCREAMING_SNAKE_CASE:				int           =   [self.normalize(image=_lowerCAmelCase			, mean=_lowerCAmelCase			, std=_lowerCAmelCase					) for image in images]

										__SCREAMING_SNAKE_CASE:				str           =   [to_channel_dimension_format(_lowerCAmelCase			, _lowerCAmelCase					) for image in images]

										__SCREAMING_SNAKE_CASE:				List[str]           =   {'''pixel_values''': images}
										return BatchFeature(data=_lowerCAmelCase			, tensor_type=_lowerCAmelCase					)



