

'''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  UpperCAmelCase :

		'''simple docstring'''






		def __init__(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_=13   ,			SCREAMING_SNAKE_CASE_=7   ,			SCREAMING_SNAKE_CASE_=True   ,			SCREAMING_SNAKE_CASE_=True   ,			SCREAMING_SNAKE_CASE_=True   ,			SCREAMING_SNAKE_CASE_=True   ,			SCREAMING_SNAKE_CASE_=99   ,			SCREAMING_SNAKE_CASE_=32   ,			SCREAMING_SNAKE_CASE_=5   ,			SCREAMING_SNAKE_CASE_=4   ,			SCREAMING_SNAKE_CASE_=37   ,			SCREAMING_SNAKE_CASE_="gelu"   ,			SCREAMING_SNAKE_CASE_=0.1   ,			SCREAMING_SNAKE_CASE_=0.1   ,			SCREAMING_SNAKE_CASE_=128   ,			SCREAMING_SNAKE_CASE_=32   ,			SCREAMING_SNAKE_CASE_=16   ,			SCREAMING_SNAKE_CASE_=2   ,			SCREAMING_SNAKE_CASE_=0.02   ,			SCREAMING_SNAKE_CASE_=3   ,			SCREAMING_SNAKE_CASE_=4   ,			SCREAMING_SNAKE_CASE_=None   ,			)  ->     Tuple:






							'''simple docstring'''



							lowerCamelCase_  				=      parent
							lowerCamelCase_  				=      batch_size
							lowerCamelCase_  				=      seq_length
							lowerCamelCase_  				=      is_training
							lowerCamelCase_  				=      use_input_mask
							lowerCamelCase_  				=      use_token_type_ids
							lowerCamelCase_  				=      use_labels
							lowerCamelCase_  				=      vocab_size
							lowerCamelCase_  				=      hidden_size
							lowerCamelCase_  				=      num_hidden_layers
							lowerCamelCase_  				=      num_attention_heads
							lowerCamelCase_  				=      intermediate_size
							lowerCamelCase_  				=      hidden_act
							lowerCamelCase_  				=      hidden_dropout_prob
							lowerCamelCase_  				=      attention_probs_dropout_prob
							lowerCamelCase_  				=      max_position_embeddings
							lowerCamelCase_  				=      type_vocab_size
							lowerCamelCase_  				=      type_sequence_label_size
							lowerCamelCase_  				=      initializer_range
							lowerCamelCase_  				=      num_labels
							lowerCamelCase_  				=      num_choices
							lowerCamelCase_  				=      scope
		def 			UpperCamelCase(					self		)  ->     Optional[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      ids_tensor([self.batch_size, self.seq_length]   ,			self.vocab_size		)

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

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

							lowerCamelCase_  				=      None
							lowerCamelCase_  				=      None
							lowerCamelCase_  				=      None
							if self.use_labels:
												lowerCamelCase_  				=      ids_tensor([self.batch_size]   ,			self.type_sequence_label_size		)
												lowerCamelCase_  				=      ids_tensor([self.batch_size, self.seq_length]   ,			self.num_labels		)
												lowerCamelCase_  				=      ids_tensor([self.batch_size]   ,			self.num_choices		)

							lowerCamelCase_  				=      self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
		def 			UpperCamelCase(					self		)  ->     Dict:






							'''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=SCREAMING_SNAKE_CASE_   ,			initializer_range=self.initializer_range   ,			)
		def 			UpperCamelCase(					self		)  ->     Optional[int]:






							'''simple docstring'''



							(
							    (
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,
							)  				=      self.prepare_config_and_inputs()

							lowerCamelCase_  				=      True
							lowerCamelCase_  				=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
							lowerCamelCase_  				=      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   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     int:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaModel(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_		)
							self.parent.assertEqual(result.last_hidden_state.shape   ,			(self.batch_size, self.seq_length, self.hidden_size)		)
							self.parent.assertEqual(result.pooler_output.shape   ,			(self.batch_size, self.hidden_size)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			)  ->     Dict:






							'''simple docstring'''



							lowerCamelCase_  				=      True
							lowerCamelCase_  				=      NezhaModel(SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(
							    SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			encoder_hidden_states=SCREAMING_SNAKE_CASE_   ,			encoder_attention_mask=SCREAMING_SNAKE_CASE_   ,			)
							lowerCamelCase_  				=      model(
							    SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			encoder_hidden_states=SCREAMING_SNAKE_CASE_   ,			)
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_		)
							self.parent.assertEqual(result.last_hidden_state.shape   ,			(self.batch_size, self.seq_length, self.hidden_size)		)
							self.parent.assertEqual(result.pooler_output.shape   ,			(self.batch_size, self.hidden_size)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Optional[int]:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaForMaskedLM(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_		)
							self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.seq_length, self.vocab_size)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Optional[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaForNextSentencePrediction(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(
							    SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_   ,			)
							self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, 2)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Optional[int]:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaForPreTraining(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(
							    SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_   ,			next_sentence_label=SCREAMING_SNAKE_CASE_   ,			)
							self.parent.assertEqual(result.prediction_logits.shape   ,			(self.batch_size, self.seq_length, self.vocab_size)		)
							self.parent.assertEqual(result.seq_relationship_logits.shape   ,			(self.batch_size, 2)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     List[str]:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaForQuestionAnswering(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(
							    SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			start_positions=SCREAMING_SNAKE_CASE_   ,			end_positions=SCREAMING_SNAKE_CASE_   ,			)
							self.parent.assertEqual(result.start_logits.shape   ,			(self.batch_size, self.seq_length)		)
							self.parent.assertEqual(result.end_logits.shape   ,			(self.batch_size, self.seq_length)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     str:






							'''simple docstring'''



							lowerCamelCase_  				=      self.num_labels
							lowerCamelCase_  				=      NezhaForSequenceClassification(SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_		)
							self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.num_labels)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     List[str]:






							'''simple docstring'''



							lowerCamelCase_  				=      self.num_labels
							lowerCamelCase_  				=      NezhaForTokenClassification(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_		)
							self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.seq_length, self.num_labels)		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     int:






							'''simple docstring'''



							lowerCamelCase_  				=      self.num_choices
							lowerCamelCase_  				=      NezhaForMultipleChoice(config=SCREAMING_SNAKE_CASE_		)
							model.to(SCREAMING_SNAKE_CASE_		)
							model.eval()
							lowerCamelCase_  				=      input_ids.unsqueeze(1		).expand(-1   ,			self.num_choices   ,			-1		).contiguous()
							lowerCamelCase_  				=      token_type_ids.unsqueeze(1		).expand(-1   ,			self.num_choices   ,			-1		).contiguous()
							lowerCamelCase_  				=      input_mask.unsqueeze(1		).expand(-1   ,			self.num_choices   ,			-1		).contiguous()
							lowerCamelCase_  				=      model(
							    SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_   ,			token_type_ids=SCREAMING_SNAKE_CASE_   ,			labels=SCREAMING_SNAKE_CASE_   ,			)
							self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.num_choices)		)






		def 			UpperCamelCase(					self		)  ->     int:






							'''simple docstring'''



							lowerCamelCase_  				=      self.prepare_config_and_inputs()
							(
							    (
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,
							)  				=      config_and_inputs
							lowerCamelCase_  				=      {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
							return config, inputs_dict


@require_torch
class  UpperCAmelCase (							UpperCAmelCase__  , UpperCAmelCase__  , UpperCAmelCase__  , unittest.TestCase    ):

		'''simple docstring'''






		SCREAMING_SNAKE_CASE_           =					(
		    (
		        NezhaModel,
		        NezhaForMaskedLM,
		        NezhaForMultipleChoice,
		        NezhaForNextSentencePrediction,
		        NezhaForPreTraining,
		        NezhaForQuestionAnswering,
		        NezhaForSequenceClassification,
		        NezhaForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		SCREAMING_SNAKE_CASE_           =					(
		    {
		        '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_           =					True
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_=False		)  ->     List[str]:






							'''simple docstring'''



							lowerCamelCase_  				=      super()._prepare_for_class(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			return_labels=SCREAMING_SNAKE_CASE_		)

							if return_labels:
												if model_class in get_values(SCREAMING_SNAKE_CASE_		):
																	lowerCamelCase_  				=      torch.zeros(
																	    (self.model_tester.batch_size, self.model_tester.seq_length)   ,			dtype=torch.long   ,			device=SCREAMING_SNAKE_CASE_		)
																	lowerCamelCase_  				=      torch.zeros(
																	    self.model_tester.batch_size   ,			dtype=torch.long   ,			device=SCREAMING_SNAKE_CASE_		)
							return inputs_dict
		def 			UpperCamelCase(					self		)  ->     List[str]:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaModelTester(self		)
							lowerCamelCase_  				=      ConfigTester(self   ,			config_class=SCREAMING_SNAKE_CASE_   ,			hidden_size=37		)
		def 			UpperCamelCase(					self		)  ->     Tuple:






							'''simple docstring'''



							self.config_tester.run_common_tests()
		def 			UpperCamelCase(					self		)  ->     Optional[int]:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     int:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     Union[str, Any]:






							'''simple docstring'''



							(
							    (
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,(
							    lowerCamelCase_
							)				,
							)  				=      self.model_tester.prepare_config_and_inputs_for_decoder()

							lowerCamelCase_  				=      None

							self.model_tester.create_and_check_model_as_decoder(
							    SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			)
		def 			UpperCamelCase(					self		)  ->     Dict:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     Tuple:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     Optional[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_next_sequence_prediction(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     Optional[int]:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     str:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     List[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_		)
		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							lowerCamelCase_  				=      self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_		)
		@slow
		def 			UpperCamelCase(					self		)  ->     Tuple:






							'''simple docstring'''



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






		@slow
		@require_torch_gpu
		def 			UpperCamelCase(					self		)  ->     List[Any]:






							'''simple docstring'''



							lowerCamelCase_				,lowerCamelCase_  				=      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_  				=      True
												lowerCamelCase_  				=      model_class(config=SCREAMING_SNAKE_CASE_		)

												lowerCamelCase_  				=      self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)
												lowerCamelCase_  				=      torch.jit.trace(
												    SCREAMING_SNAKE_CASE_   ,			(inputs_dict['input_ids'].to('cpu'		), inputs_dict['attention_mask'].to('cpu'		))		)

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


@require_torch
class  UpperCAmelCase (							unittest.TestCase    ):

		'''simple docstring'''






		@slow
		def 			UpperCamelCase(					self		)  ->     Union[str, Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaModel.from_pretrained('sijunhe/nezha-cn-base'		)
							lowerCamelCase_  				=      torch.tensor([[0, 1, 2, 3, 4, 5]]		)
							lowerCamelCase_  				=      torch.tensor([[0, 1, 1, 1, 1, 1]]		)
							with torch.no_grad():
												lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_		)[0]
							lowerCamelCase_  				=      torch.Size((1, 6, 768)		)
							self.assertEqual(output.shape   ,			SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      torch.tensor([[[0.0_685, 0.2_441, 0.1_102], [0.0_600, 0.1_906, 0.1_349], [0.0_221, 0.0_819, 0.0_586]]]		)

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






		@slow
		def 			UpperCamelCase(					self		)  ->     Optional[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      NezhaForMaskedLM.from_pretrained('sijunhe/nezha-cn-base'		)
							lowerCamelCase_  				=      torch.tensor([[0, 1, 2, 3, 4, 5]]		)
							lowerCamelCase_  				=      torch.tensor([[1, 1, 1, 1, 1, 1]]		)
							with torch.no_grad():
												lowerCamelCase_  				=      model(SCREAMING_SNAKE_CASE_   ,			attention_mask=SCREAMING_SNAKE_CASE_		)[0]
							lowerCamelCase_  				=      torch.Size((1, 6, 21128)		)
							self.assertEqual(output.shape   ,			SCREAMING_SNAKE_CASE_		)
							lowerCamelCase_  				=      torch.tensor(
							    [[-2.7_939, -1.7_902, -2.2_189], [-2.8_585, -1.8_908, -2.3_723], [-2.6_499, -1.7_750, -2.2_558]]		)

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





"""simple docstring"""
import torch

from diffusers import DDIMParallelScheduler

from .test_schedulers import SchedulerCommonTest
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							_snake_case						      =		(DDIMParallelScheduler,)
							_snake_case						      =		(('''eta''', 0.0), ('''num_inference_steps''', 5_0))




							def    A__						( self    ,				**snake_case_				)       ->  Optional[Any]:
													__lowerCAmelCase     =  {
													    """num_train_timesteps""": 1_000,
													    """beta_start""": 0.0_001,
													    """beta_end""": 0.02,
													    """beta_schedule""": """linear""",
													    """clip_sample""": True,
													}

													config.update(**snake_case_				)
													return config




							def    A__						( self    ,				**snake_case_				)       ->  Tuple:
													__lowerCAmelCase     =  self.scheduler_classes[0]
													__lowerCAmelCase     =  self.get_scheduler_config(**snake_case_				)
													__lowerCAmelCase     =  scheduler_class(**snake_case_				)

													__lowerCAmelCase   ,     __lowerCAmelCase     =  10, 0.0

													__lowerCAmelCase     =  self.dummy_model()
													__lowerCAmelCase     =  self.dummy_sample_deter

													scheduler.set_timesteps(snake_case_				)

													for t in scheduler.timesteps:
																			__lowerCAmelCase     =  model(snake_case_    ,				snake_case_				)
																			__lowerCAmelCase     =  scheduler.step(snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				).prev_sample

													return sample




							def    A__						( self				)       ->  List[str]:
													for timesteps in [100, 500, 1_000]:
																			self.check_over_configs(num_train_timesteps=snake_case_				)




							def    A__						( self				)       ->  Tuple:
													for steps_offset in [0, 1]:
																			self.check_over_configs(steps_offset=snake_case_				)

													__lowerCAmelCase     =  self.scheduler_classes[0]
													__lowerCAmelCase     =  self.get_scheduler_config(steps_offset=1				)
													__lowerCAmelCase     =  scheduler_class(**snake_case_				)
													scheduler.set_timesteps(5				)
													assert torch.equal(scheduler.timesteps    ,				torch.LongTensor([801, 601, 401, 201, 1]				)				)




							def    A__						( self				)       ->  int:
													for beta_start, beta_end in zip([0.0_001, 0.001, 0.01, 0.1]    ,				[0.002, 0.02, 0.2, 2]				):
																			self.check_over_configs(beta_start=snake_case_    ,				beta_end=snake_case_				)




							def    A__						( self				)       ->  str:
													for schedule in ["linear", "squaredcos_cap_v2"]:
																			self.check_over_configs(beta_schedule=snake_case_				)




							def    A__						( self				)       ->  Tuple:
													for prediction_type in ["epsilon", "v_prediction"]:
																			self.check_over_configs(prediction_type=snake_case_				)




							def    A__						( self				)       ->  Optional[int]:
													for clip_sample in [True, False]:
																			self.check_over_configs(clip_sample=snake_case_				)




							def    A__						( self				)       ->  Dict:
													for timestep_spacing in ["trailing", "leading"]:
																			self.check_over_configs(timestep_spacing=snake_case_				)




							def    A__						( self				)       ->  str:
													for rescale_betas_zero_snr in [True, False]:
																			self.check_over_configs(rescale_betas_zero_snr=snake_case_				)




							def    A__						( self				)       ->  Optional[int]:
													self.check_over_configs(thresholding=snake_case_				)
													for threshold in [0.5, 1.0, 2.0]:
																			for prediction_type in ["epsilon", "v_prediction"]:
																									self.check_over_configs(
																									    thresholding=snake_case_    ,				prediction_type=snake_case_    ,				sample_max_value=snake_case_    ,				)




							def    A__						( self				)       ->  Union[str, Any]:
													for t in [1, 10, 49]:
																			self.check_over_forward(time_step=snake_case_				)




							def    A__						( self				)       ->  Optional[int]:
													for t, num_inference_steps in zip([1, 10, 50]    ,				[10, 50, 500]				):
																			self.check_over_forward(time_step=snake_case_    ,				num_inference_steps=snake_case_				)




							def    A__						( self				)       ->  Optional[int]:
													for t, eta in zip([1, 10, 49]    ,				[0.0, 0.5, 1.0]				):
																			self.check_over_forward(time_step=snake_case_    ,				eta=snake_case_				)




							def    A__						( self				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.scheduler_classes[0]
													__lowerCAmelCase     =  self.get_scheduler_config()
													__lowerCAmelCase     =  scheduler_class(**snake_case_				)

													assert torch.sum(torch.abs(scheduler._get_variance(0    ,				0				) - 0.0				)				) < 1e-5
													assert torch.sum(torch.abs(scheduler._get_variance(420    ,				400				) - 0.14_771				)				) < 1e-5
													assert torch.sum(torch.abs(scheduler._get_variance(980    ,				960				) - 0.32_460				)				) < 1e-5
													assert torch.sum(torch.abs(scheduler._get_variance(0    ,				0				) - 0.0				)				) < 1e-5
													assert torch.sum(torch.abs(scheduler._get_variance(487    ,				486				) - 0.00_979				)				) < 1e-5
													assert torch.sum(torch.abs(scheduler._get_variance(999    ,				998				) - 0.02				)				) < 1e-5




							def    A__						( self				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.scheduler_classes[0]
													__lowerCAmelCase     =  self.get_scheduler_config()
													__lowerCAmelCase     =  scheduler_class(**snake_case_				)

													__lowerCAmelCase   ,     __lowerCAmelCase     =  10, 0.0
													scheduler.set_timesteps(snake_case_				)

													__lowerCAmelCase     =  self.dummy_model()
													__lowerCAmelCase     =  self.dummy_sample_deter
													__lowerCAmelCase     =  self.dummy_sample_deter + 0.1
													__lowerCAmelCase     =  self.dummy_sample_deter - 0.1

													__lowerCAmelCase     =  samplea.shape[0]
													__lowerCAmelCase     =  torch.stack([samplea, samplea, samplea]    ,				dim=0				)
													__lowerCAmelCase     =  torch.arange(snake_case_				)[0:3, None].repeat(1    ,				snake_case_				)

													__lowerCAmelCase     =  model(samples.flatten(0    ,				1				)    ,				timesteps.flatten(0    ,				1				)				)
													__lowerCAmelCase     =  scheduler.batch_step_no_noise(snake_case_    ,				timesteps.flatten(0    ,				1				)    ,				samples.flatten(0    ,				1				)    ,				snake_case_				)

													__lowerCAmelCase     =  torch.sum(torch.abs(snake_case_				)				)
													__lowerCAmelCase     =  torch.mean(torch.abs(snake_case_				)				)

													assert abs(result_sum.item() - 1_147.7_904				) < 1e-2
													assert abs(result_mean.item() - 0.4_982				) < 1e-3




							def    A__						( self				)       ->  Dict:
													__lowerCAmelCase     =  self.full_loop()

													__lowerCAmelCase     =  torch.sum(torch.abs(snake_case_				)				)
													__lowerCAmelCase     =  torch.mean(torch.abs(snake_case_				)				)

													assert abs(result_sum.item() - 172.0_067				) < 1e-2
													assert abs(result_mean.item() - 0.223_967				) < 1e-3




							def    A__						( self				)       ->  str:
													__lowerCAmelCase     =  self.full_loop(prediction_type="""v_prediction"""				)

													__lowerCAmelCase     =  torch.sum(torch.abs(snake_case_				)				)
													__lowerCAmelCase     =  torch.mean(torch.abs(snake_case_				)				)

													assert abs(result_sum.item() - 52.5_302				) < 1e-2
													assert abs(result_mean.item() - 0.0_684				) < 1e-3




							def    A__						( self				)       ->  Tuple:
													# We specify different beta, so that the first alpha is 0.99
													__lowerCAmelCase     =  self.full_loop(set_alpha_to_one=snake_case_    ,				beta_start=0.01				)
													__lowerCAmelCase     =  torch.sum(torch.abs(snake_case_				)				)
													__lowerCAmelCase     =  torch.mean(torch.abs(snake_case_				)				)

													assert abs(result_sum.item() - 149.8_295				) < 1e-2
													assert abs(result_mean.item() - 0.1_951				) < 1e-3




							def    A__						( self				)       ->  Union[str, Any]:
													# We specify different beta, so that the first alpha is 0.99
													__lowerCAmelCase     =  self.full_loop(set_alpha_to_one=snake_case_    ,				beta_start=0.01				)
													__lowerCAmelCase     =  torch.sum(torch.abs(snake_case_				)				)
													__lowerCAmelCase     =  torch.mean(torch.abs(snake_case_				)				)

													assert abs(result_sum.item() - 149.0_784				) < 1e-2
													assert abs(result_mean.item() - 0.1_941				) < 1e-3
