



'''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 	lowerCAmelCase_:
		'''simple docstring'''




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




		def 							UpperCAmelCase_     (   self				)   ->		str:
						lowerCAmelCase__			:  Tuple											= ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size				)

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

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

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

						lowerCAmelCase__			:  List[str]											= self.get_config()

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




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




		def 							UpperCAmelCase_     (   self				)   ->		Any:
						(
						    lowerCAmelCase__ 
						)			:  int											= self.prepare_config_and_inputs()

						lowerCAmelCase__			:  str											= True
						lowerCAmelCase__			:  str											= 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  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Optional[Any]:
						lowerCAmelCase__			:  Any											= NezhaModel(config=_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  Dict											= model(_a  ,attention_mask=_a  ,token_type_ids=_a				)
						lowerCAmelCase__			:  Optional[int]											= model(_a  ,token_type_ids=_a				)
						lowerCAmelCase__			:  List[Any]											= model(_a				)
						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  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,)   ->		Optional[Any]:
						lowerCAmelCase__			:  List[Any]											= True
						lowerCAmelCase__			:  Optional[int]											= NezhaModel(_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  Tuple											= model(
						    _a  ,attention_mask=_a  ,token_type_ids=_a  ,encoder_hidden_states=_a  ,encoder_attention_mask=_a  ,)
						lowerCAmelCase__			:  Union[str, Any]											= model(
						    _a  ,attention_mask=_a  ,token_type_ids=_a  ,encoder_hidden_states=_a  ,)
						lowerCAmelCase__			:  Any											= model(_a  ,attention_mask=_a  ,token_type_ids=_a				)
						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  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		List[str]:
						lowerCAmelCase__			:  Optional[Any]											= NezhaForMaskedLM(config=_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  str											= model(_a  ,attention_mask=_a  ,token_type_ids=_a  ,labels=_a				)
						self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.vocab_size)				)




		def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Tuple:
						lowerCAmelCase__			:  List[Any]											= NezhaForNextSentencePrediction(config=_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  Dict											= model(
						    _a  ,attention_mask=_a  ,token_type_ids=_a  ,labels=_a  ,)
						self.parent.assertEqual(result.logits.shape  ,(self.batch_size, 2)				)




		def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Optional[int]:
						lowerCAmelCase__			:  Any											= NezhaForPreTraining(config=_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  Dict											= model(
						    _a  ,attention_mask=_a  ,token_type_ids=_a  ,labels=_a  ,next_sentence_label=_a  ,)
						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  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Optional[int]:
						lowerCAmelCase__			:  Optional[Any]											= NezhaForQuestionAnswering(config=_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  Dict											= model(
						    _a  ,attention_mask=_a  ,token_type_ids=_a  ,start_positions=_a  ,end_positions=_a  ,)
						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  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		str:
						lowerCAmelCase__			:  int											= self.num_labels
						lowerCAmelCase__			:  List[str]											= NezhaForSequenceClassification(_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  List[str]											= model(_a  ,attention_mask=_a  ,token_type_ids=_a  ,labels=_a				)
						self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.num_labels)				)




		def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Tuple:
						lowerCAmelCase__			:  List[Any]											= self.num_labels
						lowerCAmelCase__			:  Any											= NezhaForTokenClassification(config=_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  List[Any]											= model(_a  ,attention_mask=_a  ,token_type_ids=_a  ,labels=_a				)
						self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.num_labels)				)




		def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		List[str]:
						lowerCAmelCase__			:  List[str]											= self.num_choices
						lowerCAmelCase__			:  Tuple											= NezhaForMultipleChoice(config=_a				)
						model.to(_a				)
						model.eval()
						lowerCAmelCase__			:  Optional[Any]											= input_ids.unsqueeze(1				).expand(-1  ,self.num_choices  ,-1				).contiguous()
						lowerCAmelCase__			:  Tuple											= token_type_ids.unsqueeze(1				).expand(-1  ,self.num_choices  ,-1				).contiguous()
						lowerCAmelCase__			:  Any											= input_mask.unsqueeze(1				).expand(-1  ,self.num_choices  ,-1				).contiguous()
						lowerCAmelCase__			:  Tuple											= model(
						    _a  ,attention_mask=_a  ,token_type_ids=_a  ,labels=_a  ,)
						self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.num_choices)				)






		def 							UpperCAmelCase_     (   self				)   ->		List[str]:
						lowerCAmelCase__			:  str											= self.prepare_config_and_inputs()
						(
						    lowerCAmelCase__ 
						)			:  List[str]											= config_and_inputs
						lowerCAmelCase__			:  List[Any]											= {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
						return config, inputs_dict







@require_torch
class 	lowerCAmelCase_(    __lowercase						,			__lowercase						,			__lowercase						,			unittest.TestCase					):
		'''simple docstring'''




		__lowercase							:				Tuple           =				(
		    (
		        NezhaModel,
		        NezhaForMaskedLM,
		        NezhaForMultipleChoice,
		        NezhaForNextSentencePrediction,
		        NezhaForPreTraining,
		        NezhaForQuestionAnswering,
		        NezhaForSequenceClassification,
		        NezhaForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		__lowercase							:				Optional[int]           =				(
		    {
		        "feature-extraction": NezhaModel,
		        "fill-mask": NezhaForMaskedLM,
		        "question-answering": NezhaForQuestionAnswering,
		        "text-classification": NezhaForSequenceClassification,
		        "token-classification": NezhaForTokenClassification,
		        "zero-shot": NezhaForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		__lowercase							:				Dict           =				True




		def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase=False				)   ->		List[Any]:
						lowerCAmelCase__			:  Any											= super()._prepare_for_class(_a  ,_a  ,return_labels=_a				)

						if return_labels:
										if model_class in get_values(_a				):
														lowerCAmelCase__			:  Dict											= torch.zeros(
														    (self.model_tester.batch_size, self.model_tester.seq_length)  ,dtype=torch.long  ,device=_a				)
														lowerCAmelCase__			:  int											= torch.zeros(
														    self.model_tester.batch_size  ,dtype=torch.long  ,device=_a				)
						return inputs_dict




		def 							UpperCAmelCase_     (   self				)   ->		str:
						lowerCAmelCase__			:  str											= NezhaModelTester(self				)
						lowerCAmelCase__			:  Optional[Any]											= ConfigTester(self  ,config_class=_a  ,hidden_size=37				)




		def 							UpperCAmelCase_     (   self				)   ->		Optional[Any]:
						self.config_tester.run_common_tests()




		def 							UpperCAmelCase_     (   self				)   ->		Any:
						lowerCAmelCase__			:  int											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		Any:
						lowerCAmelCase__			:  str											= self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		Union[str, Any]:
						# This regression test was failing with PyTorch < 1.3
						(
						    lowerCAmelCase__ 
						)			:  List[str]											= self.model_tester.prepare_config_and_inputs_for_decoder()

						lowerCAmelCase__			:  int											= None

						self.model_tester.create_and_check_model_as_decoder(
						    _a  ,_a  ,_a  ,_a  ,_a  ,_a  ,_a  ,_a  ,_a  ,)




		def 							UpperCAmelCase_     (   self				)   ->		Tuple:
						lowerCAmelCase__			:  int											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		List[str]:
						lowerCAmelCase__			:  str											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		Union[str, Any]:
						lowerCAmelCase__			:  Union[str, Any]											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_next_sequence_prediction(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		Tuple:
						lowerCAmelCase__			:  Union[str, Any]											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		Union[str, Any]:
						lowerCAmelCase__			:  Any											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		Dict:
						lowerCAmelCase__			:  int											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*_a				)




		def 							UpperCAmelCase_     (   self				)   ->		Optional[Any]:
						lowerCAmelCase__			:  str											= self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*_a				)




		@slow
		def 							UpperCAmelCase_     (   self				)   ->		Tuple:
						for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__			:  Dict											= NezhaModel.from_pretrained(_a				)
										self.assertIsNotNone(_a				)






		@slow
		@require_torch_gpu
		def 							UpperCAmelCase_     (   self				)   ->		Optional[Any]:
						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__			:  List[Any]											= True
										lowerCAmelCase__			:  List[Any]											= model_class(config=_a				)

										lowerCAmelCase__			:  Optional[int]											= self._prepare_for_class(_a  ,_a				)
										lowerCAmelCase__			:  Tuple											= torch.jit.trace(
										    _a  ,(inputs_dict["""input_ids"""].to("""cpu"""				), inputs_dict["""attention_mask"""].to("""cpu"""				))				)

										with tempfile.TemporaryDirectory() as tmp:
														torch.jit.save(_a  ,os.path.join(_a  ,"""bert.pt"""				)				)
														lowerCAmelCase__			:  Optional[Any]											= torch.jit.load(os.path.join(_a  ,"""bert.pt"""				)  ,map_location=_a				)
														loaded(inputs_dict["""input_ids"""].to(_a				)  ,inputs_dict["""attention_mask"""].to(_a				)				)







@require_torch
class 	lowerCAmelCase_(    unittest.TestCase					):
		'''simple docstring'''




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






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







'''simple docstring'''





def        _SCREAMING_SNAKE_CASE  (	UpperCamelCase = 1000      ):



    """simple docstring"""





    lowerCAmelCase__			:  Union[str, Any]											= -1
    lowerCAmelCase__			:  Optional[Any]											= 0
    for a in range(1							,   n // 3      ):
        # Solving the two equations a**2+b**2=c**2 and a+b+c=N eliminating c
        lowerCAmelCase__			:  Optional[Any]											= (n * n - 2 * a * n) // (2 * n - 2 * a)
        lowerCAmelCase__			:  Tuple											= n - a - b
        if c * c == (a * a + b * b):
            lowerCAmelCase__			:  int											= a * b * c
            if candidate >= product:
                lowerCAmelCase__			:  Any											= candidate
    return product


if __name__ == "__main__":
 print(F"""{solution() = }""")


