

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				,       snake_case				,       snake_case=13				,       snake_case=7				,       snake_case=True				,       snake_case=True				,       snake_case=True				,       snake_case=True				,       snake_case=99				,       snake_case=32				,       snake_case=5				,       snake_case=4				,       snake_case=37				,       snake_case="gelu"				,       snake_case=0.1				,       snake_case=0.1				,       snake_case=128				,       snake_case=32				,       snake_case=16				,       snake_case=2				,       snake_case=0.02				,       snake_case=3				,       snake_case=4				,       snake_case=None				,       ):
											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    a							( self	):
											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    a							( 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    a							( self	):
											(
											    (
											    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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   NezhaModel(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case	)
											snake_case_          =   model(snake_case				,       token_type_ids=snake_case	)
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       ):
											snake_case_          =   True
											snake_case_          =   NezhaModel(snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   model(
											    snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       encoder_hidden_states=snake_case				,       encoder_attention_mask=snake_case				,       )
											snake_case_          =   model(
											    snake_case				,       attention_mask=snake_case				,       token_type_ids=snake_case				,       encoder_hidden_states=snake_case				,       )
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   NezhaForMaskedLM(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   NezhaForNextSentencePrediction(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   NezhaForPreTraining(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   NezhaForQuestionAnswering(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   self.num_labels
											snake_case_          =   NezhaForSequenceClassification(snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   self.num_labels
											snake_case_          =   NezhaForTokenClassification(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   self.num_choices
											snake_case_          =   NezhaForMultipleChoice(config=snake_case	)
											model.to(snake_case	)
											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(
											    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    a							( self	):
											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 			lowercase    (	lowercase_     ,  lowercase_     ,  lowercase_     ,  unittest.TestCase      ):
				__SCREAMING_SNAKE_CASE :			Optional[int]	      =      (
				    (
				        NezhaModel,
				        NezhaForMaskedLM,
				        NezhaForMultipleChoice,
				        NezhaForNextSentencePrediction,
				        NezhaForPreTraining,
				        NezhaForQuestionAnswering,
				        NezhaForSequenceClassification,
				        NezhaForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				__SCREAMING_SNAKE_CASE :			Union[str, Any]	      =      (
				    {
				        '''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 :			List[Any]	      =      True




				def    a							( self				,       snake_case				,       snake_case				,       snake_case=False	):
											snake_case_          =   super()._prepare_for_class(snake_case				,       snake_case				,       return_labels=snake_case	)

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




				def    a							( self	):
											snake_case_          =   NezhaModelTester(self	)
											snake_case_          =   ConfigTester(self				,       config_class=snake_case				,       hidden_size=37	)




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




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*snake_case	)




				def    a							( self	):
											# This regression test was failing with PyTorch < 1.3
											(
											    (
											    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(
											    snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       )




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_next_sequence_prediction(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_pretraining(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*snake_case	)




				@slow
				def    a							( self	):
											for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		snake_case_          =   NezhaModel.from_pretrained(snake_case	)
																		self.assertIsNotNone(snake_case	)




				@slow
				@require_torch_gpu
				def    a							( self	):
											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=snake_case	)

																		snake_case_          =   self._prepare_for_class(snake_case				,       snake_case	)
																		snake_case_          =   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'	)	)
																									snake_case_          =   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    a							( self	):
											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(snake_case				,       attention_mask=snake_case	)[0]
											snake_case_          =   torch.Size((1, 6, 768)	)
											self.assertEqual(output.shape				,       snake_case	)
											snake_case_          =   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    a							( self	):
											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(snake_case				,       attention_mask=snake_case	)[0]
											snake_case_          =   torch.Size((1, 6, 2_1128)	)
											self.assertEqual(output.shape				,       snake_case	)
											snake_case_          =   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	)	)







from __future__ import annotations





def  __lowerCamelCase     (			UpperCamelCase__       ,							UpperCamelCase__ = None       ,							UpperCamelCase__ = None				):







							'''simple docstring'''






							if start is None:
														snake_case_          =   0

							if end is None:
														snake_case_          =   len(UpperCamelCase__				) - 1

							if start >= end:
														return

							snake_case_          =   (start + end) // 2

							slowsort(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__				)
							slowsort(UpperCamelCase__       ,							mid + 1       ,							UpperCamelCase__				)

							if sequence[end] < sequence[mid]:
														snake_case_    ,   snake_case_          =   sequence[mid], sequence[end]

							slowsort(UpperCamelCase__       ,							UpperCamelCase__       ,							end - 1				)


if __name__ == "__main__":
					from doctest import testmod

					testmod()




