
'''simple docstring'''


import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    MODEL_FOR_PRETRAINING_MAPPING,
					    NezhaForMaskedLM,
					    NezhaForMultipleChoice,
					    NezhaForNextSentencePrediction,
					    NezhaForPreTraining,
					    NezhaForQuestionAnswering,
					    NezhaForSequenceClassification,
					    NezhaForTokenClassification,
					    NezhaModel,
					)
					from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST

class 		__lowercase :

			def __init__(self	,      A	,      A=1_3	,      A=7	,      A=True	,      A=True	,      A=True	,      A=True	,      A=9_9	,      A=3_2	,      A=5	,      A=4	,      A=3_7	,      A="gelu"	,      A=0.1	,      A=0.1	,      A=1_2_8	,      A=3_2	,      A=1_6	,      A=2	,      A=0.02	,      A=3	,      A=4	,      A=None	,      ):
								lowerCamelCase_    :					Dict   =						parent
								lowerCamelCase_    :					str   =						batch_size
								lowerCamelCase_    :					Dict   =						seq_length
								lowerCamelCase_    :					List[str]   =						is_training
								lowerCamelCase_    :					List[Any]   =						use_input_mask
								lowerCamelCase_    :					List[Any]   =						use_token_type_ids
								lowerCamelCase_    :					Tuple   =						use_labels
								lowerCamelCase_    :					Union[str, Any]   =						vocab_size
								lowerCamelCase_    :					Dict   =						hidden_size
								lowerCamelCase_    :					List[Any]   =						num_hidden_layers
								lowerCamelCase_    :					int   =						num_attention_heads
								lowerCamelCase_    :					Tuple   =						intermediate_size
								lowerCamelCase_    :					int   =						hidden_act
								lowerCamelCase_    :					Tuple   =						hidden_dropout_prob
								lowerCamelCase_    :					Tuple   =						attention_probs_dropout_prob
								lowerCamelCase_    :					Optional[int]   =						max_position_embeddings
								lowerCamelCase_    :					Optional[Any]   =						type_vocab_size
								lowerCamelCase_    :					Dict   =						type_sequence_label_size
								lowerCamelCase_    :					Any   =						initializer_range
								lowerCamelCase_    :					int   =						num_labels
								lowerCamelCase_    :					Union[str, Any]   =						num_choices
								lowerCamelCase_    :					str   =						scope

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					List[str]   =						ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size					)

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

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

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

								lowerCamelCase_    :					str   =						self.get_config()

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

			def 				UpperCAmelCase__       (self					):
								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=_snake_case	,      initializer_range=self.initializer_range	,      )

			def 				UpperCAmelCase__       (self					):
								(
								    (
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					       
								)    :					Tuple   =						self.prepare_config_and_inputs()

								lowerCamelCase_    :					Any   =						True
								lowerCamelCase_    :					Tuple   =						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
								lowerCamelCase_    :					Tuple   =						ids_tensor([self.batch_size, self.seq_length]	,      vocab_size=2					)

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    input_mask,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					Optional[Any]   =						NezhaModel(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					Union[str, Any]   =						model(_snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case					)
								lowerCamelCase_    :					Dict   =						model(_snake_case	,      token_type_ids=_snake_case					)
								lowerCamelCase_    :					str   =						model(_snake_case					)
								self.parent.assertEqual(result.last_hidden_state.shape	,      (self.batch_size, self.seq_length, self.hidden_size)					)
								self.parent.assertEqual(result.pooler_output.shape	,      (self.batch_size, self.hidden_size)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A	,      A	,      A	,      ):
								lowerCamelCase_    :					List[str]   =						True
								lowerCamelCase_    :					Optional[Any]   =						NezhaModel(_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					Any   =						model(
								    _snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      encoder_hidden_states=_snake_case	,      encoder_attention_mask=_snake_case	,      )
								lowerCamelCase_    :					Any   =						model(
								    _snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      encoder_hidden_states=_snake_case	,      )
								lowerCamelCase_    :					Union[str, Any]   =						model(_snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case					)
								self.parent.assertEqual(result.last_hidden_state.shape	,      (self.batch_size, self.seq_length, self.hidden_size)					)
								self.parent.assertEqual(result.pooler_output.shape	,      (self.batch_size, self.hidden_size)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					Tuple   =						NezhaForMaskedLM(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					Any   =						model(_snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      labels=_snake_case					)
								self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length, self.vocab_size)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					Union[str, Any]   =						NezhaForNextSentencePrediction(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					List[str]   =						model(
								    _snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      labels=_snake_case	,      )
								self.parent.assertEqual(result.logits.shape	,      (self.batch_size, 2)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					List[Any]   =						NezhaForPreTraining(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					List[Any]   =						model(
								    _snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      labels=_snake_case	,      next_sentence_label=_snake_case	,      )
								self.parent.assertEqual(result.prediction_logits.shape	,      (self.batch_size, self.seq_length, self.vocab_size)					)
								self.parent.assertEqual(result.seq_relationship_logits.shape	,      (self.batch_size, 2)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					List[str]   =						NezhaForQuestionAnswering(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					Union[str, Any]   =						model(
								    _snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      start_positions=_snake_case	,      end_positions=_snake_case	,      )
								self.parent.assertEqual(result.start_logits.shape	,      (self.batch_size, self.seq_length)					)
								self.parent.assertEqual(result.end_logits.shape	,      (self.batch_size, self.seq_length)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					List[str]   =						self.num_labels
								lowerCamelCase_    :					str   =						NezhaForSequenceClassification(_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					List[Any]   =						model(_snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      labels=_snake_case					)
								self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.num_labels)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					Optional[Any]   =						self.num_labels
								lowerCamelCase_    :					int   =						NezhaForTokenClassification(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					Optional[int]   =						model(_snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      labels=_snake_case					)
								self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length, self.num_labels)					)

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A	,      A	,      A	,      A	,      A					):
								lowerCamelCase_    :					List[Any]   =						self.num_choices
								lowerCamelCase_    :					Union[str, Any]   =						NezhaForMultipleChoice(config=_snake_case					)
								model.to(_snake_case					)
								model.eval()
								lowerCamelCase_    :					int   =						input_ids.unsqueeze(1					).expand(-1	,      self.num_choices	,      -1					).contiguous()
								lowerCamelCase_    :					Optional[Any]   =						token_type_ids.unsqueeze(1					).expand(-1	,      self.num_choices	,      -1					).contiguous()
								lowerCamelCase_    :					List[str]   =						input_mask.unsqueeze(1					).expand(-1	,      self.num_choices	,      -1					).contiguous()
								lowerCamelCase_    :					List[Any]   =						model(
								    _snake_case	,      attention_mask=_snake_case	,      token_type_ids=_snake_case	,      labels=_snake_case	,      )
								self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.num_choices)					)






			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Optional[Any]   =						self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					       
								)    :					Optional[Any]   =						config_and_inputs
								lowerCamelCase_    :					Optional[Any]   =						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
								return config, inputs_dict

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

			def 				UpperCAmelCase__       (self	,      A	,      A	,      A=False					):
								lowerCamelCase_    :					Dict   =						super()._prepare_for_class(_snake_case	,      _snake_case	,      return_labels=_snake_case					)

								if return_labels:
													if model_class in get_values(_snake_case					):
																		lowerCamelCase_    :					str   =						torch.zeros(
																		    (self.model_tester.batch_size, self.model_tester.seq_length)	,      dtype=torch.long	,      device=_snake_case					)
																		lowerCamelCase_    :					Dict   =						torch.zeros(
																		    self.model_tester.batch_size	,      dtype=torch.long	,      device=_snake_case					)
								return inputs_dict

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					List[Any]   =						NezhaModelTester(self					)
								lowerCamelCase_    :					List[Any]   =						ConfigTester(self	,      config_class=_snake_case	,      hidden_size=3_7					)

			def 				UpperCAmelCase__       (self					):
								self.config_tester.run_common_tests()

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Union[str, Any]   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Any   =						self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								(
								    (
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					(
								    lowerCamelCase_       
								),					       
								)    :					Optional[Any]   =						self.model_tester.prepare_config_and_inputs_for_decoder()

								lowerCamelCase_    :					int   =						None

								self.model_tester.create_and_check_model_as_decoder(
								    _snake_case	,      _snake_case	,      _snake_case	,      _snake_case	,      _snake_case	,      _snake_case	,      _snake_case	,      _snake_case	,      _snake_case	,      )

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Dict   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Optional[int]   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					List[str]   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_next_sequence_prediction(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Union[str, Any]   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Dict   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Union[str, Any]   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*_snake_case					)

			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					List[Any]   =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*_snake_case					)

			@slow
			def 				UpperCAmelCase__       (self					):


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






			@slow
			@require_torch_gpu
			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_,					lowerCamelCase_    :					Any   =						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

													lowerCamelCase_    :					Tuple   =						True
													lowerCamelCase_    :					Tuple   =						model_class(config=_snake_case					)

													lowerCamelCase_    :					List[Any]   =						self._prepare_for_class(_snake_case	,      _snake_case					)
													lowerCamelCase_    :					Optional[int]   =						torch.jit.trace(
													    _snake_case	,      (inputs_dict['''input_ids'''].to('''cpu'''					), inputs_dict['''attention_mask'''].to('''cpu'''					))					)

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

@require_torch
class 		__lowercase (     unittest.TestCase			):

			@slow
			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Union[str, Any]   =						NezhaModel.from_pretrained('''sijunhe/nezha-cn-base'''					)
								lowerCamelCase_    :					Dict   =						torch.tensor([[0, 1, 2, 3, 4, 5]]					)
								lowerCamelCase_    :					List[Any]   =						torch.tensor([[0, 1, 1, 1, 1, 1]]					)
								with torch.no_grad():
													lowerCamelCase_    :					int   =						model(_snake_case	,      attention_mask=_snake_case					)[0]
								lowerCamelCase_    :					List[Any]   =						torch.Size((1, 6, 7_6_8)					)
								self.assertEqual(output.shape	,      _snake_case					)
								lowerCamelCase_    :					int   =						torch.tensor([[[0.06_85, 0.24_41, 0.11_02], [0.06_00, 0.19_06, 0.13_49], [0.02_21, 0.08_19, 0.05_86]]]					)

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






			@slow
			def 				UpperCAmelCase__       (self					):
								lowerCamelCase_    :					Optional[int]   =						NezhaForMaskedLM.from_pretrained('''sijunhe/nezha-cn-base'''					)
								lowerCamelCase_    :					List[str]   =						torch.tensor([[0, 1, 2, 3, 4, 5]]					)
								lowerCamelCase_    :					Union[str, Any]   =						torch.tensor([[1, 1, 1, 1, 1, 1]]					)
								with torch.no_grad():
													lowerCamelCase_    :					Optional[Any]   =						model(_snake_case	,      attention_mask=_snake_case					)[0]
								lowerCamelCase_    :					Any   =						torch.Size((1, 6, 2_1_1_2_8)					)
								self.assertEqual(output.shape	,      _snake_case					)
								lowerCamelCase_    :					Optional[int]   =						torch.tensor(
								    [[-2.79_39, -1.79_02, -2.21_89], [-2.85_85, -1.89_08, -2.37_23], [-2.64_99, -1.77_50, -2.25_58]]					)

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








snake_case_ : Dict    		=       {
    "Pillow": "Pillow",
    "accelerate": "accelerate>=0.11.0",
    "compel": "compel==0.1.8",
    "black": "black~=23.1",
    "datasets": "datasets",
    "filelock": "filelock",
    "flax": "flax>=0.4.1",
    "hf-doc-builder": "hf-doc-builder>=0.3.0",
    "huggingface-hub": "huggingface-hub>=0.13.2",
    "requests-mock": "requests-mock==1.10.0",
    "importlib_metadata": "importlib_metadata",
    "invisible-watermark": "invisible-watermark",
    "isort": "isort>=5.5.4",
    "jax": "jax>=0.2.8,!=0.3.2",
    "jaxlib": "jaxlib>=0.1.65",
    "Jinja2": "Jinja2",
    "k-diffusion": "k-diffusion>=0.0.12",
    "torchsde": "torchsde",
    "note_seq": "note_seq",
    "librosa": "librosa",
    "numpy": "numpy",
    "omegaconf": "omegaconf",
    "parameterized": "parameterized",
    "protobuf": "protobuf>=3.20.3,<4",
    "pytest": "pytest",
    "pytest-timeout": "pytest-timeout",
    "pytest-xdist": "pytest-xdist",
    "ruff": "ruff>=0.0.241",
    "safetensors": "safetensors",
    "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
    "scipy": "scipy",
    "onnx": "onnx",
    "regex": "regex!=2019.12.17",
    "requests": "requests",
    "tensorboard": "tensorboard",
    "torch": "torch>=1.4",
    "torchvision": "torchvision",
    "transformers": "transformers>=4.25.1",
    "urllib3": "urllib3<=2.0.0",
}






