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   :					Optional[int]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					Any=13	,							_UpperCamelCase   :					int=7	,							_UpperCamelCase   :					Optional[Any]=True	,							_UpperCamelCase   :					Tuple=True	,							_UpperCamelCase   :					Any=True	,							_UpperCamelCase   :					Optional[Any]=True	,							_UpperCamelCase   :					Optional[int]=99	,							_UpperCamelCase   :					str=32	,							_UpperCamelCase   :					Tuple=5	,							_UpperCamelCase   :					Optional[int]=4	,							_UpperCamelCase   :					Optional[int]=37	,							_UpperCamelCase   :					List[Any]="gelu"	,							_UpperCamelCase   :					int=0.1	,							_UpperCamelCase   :					Union[str, Any]=0.1	,							_UpperCamelCase   :					int=128	,							_UpperCamelCase   :					Any=32	,							_UpperCamelCase   :					Dict=16	,							_UpperCamelCase   :					List[str]=2	,							_UpperCamelCase   :					Optional[Any]=0.0_2	,							_UpperCamelCase   :					Optional[int]=3	,							_UpperCamelCase   :					Any=4	,							_UpperCamelCase   :					int=None	,							)					->				Any:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       parent
												SCREAMING_SNAKE_CASE    			=       batch_size
												SCREAMING_SNAKE_CASE    			=       seq_length
												SCREAMING_SNAKE_CASE    			=       is_training
												SCREAMING_SNAKE_CASE    			=       use_input_mask
												SCREAMING_SNAKE_CASE    			=       use_token_type_ids
												SCREAMING_SNAKE_CASE    			=       use_labels
												SCREAMING_SNAKE_CASE    			=       vocab_size
												SCREAMING_SNAKE_CASE    			=       hidden_size
												SCREAMING_SNAKE_CASE    			=       num_hidden_layers
												SCREAMING_SNAKE_CASE    			=       num_attention_heads
												SCREAMING_SNAKE_CASE    			=       intermediate_size
												SCREAMING_SNAKE_CASE    			=       hidden_act
												SCREAMING_SNAKE_CASE    			=       hidden_dropout_prob
												SCREAMING_SNAKE_CASE    			=       attention_probs_dropout_prob
												SCREAMING_SNAKE_CASE    			=       max_position_embeddings
												SCREAMING_SNAKE_CASE    			=       type_vocab_size
												SCREAMING_SNAKE_CASE    			=       type_sequence_label_size
												SCREAMING_SNAKE_CASE    			=       initializer_range
												SCREAMING_SNAKE_CASE    			=       num_labels
												SCREAMING_SNAKE_CASE    			=       num_choices
												SCREAMING_SNAKE_CASE    			=       scope




							def 	__snake_case(						self   :					Any      )					->				Optional[int]:



												'''simple docstring'''

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

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

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

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

												SCREAMING_SNAKE_CASE    			=       self.get_config()

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




							def 	__snake_case(						self   :					List[str]      )					->				int:



												'''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=_UpperCamelCase	,							initializer_range=self.initializer_range	,							)




							def 	__snake_case(						self   :					Union[str, Any]      )					->				List[Any]:



												'''simple docstring'''

												(
												    (
												    SCREAMING_SNAKE_CASE
												)			,   (
												    SCREAMING_SNAKE_CASE
												)			,   (
												    SCREAMING_SNAKE_CASE
												)			,   (
												    SCREAMING_SNAKE_CASE
												)			,   (
												    SCREAMING_SNAKE_CASE
												)			,   (
												    SCREAMING_SNAKE_CASE
												)			,   (
												    SCREAMING_SNAKE_CASE
												)			,   
												)    			=       self.prepare_config_and_inputs()

												SCREAMING_SNAKE_CASE    			=       True
												SCREAMING_SNAKE_CASE    			=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
												SCREAMING_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 	__snake_case(						self   :					Dict	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					str	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					Optional[int]	,							_UpperCamelCase   :					Optional[Any]      )					->				str:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaModel(config=_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase      )
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							token_type_ids=_UpperCamelCase      )
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase      )
												self.parent.assertEqual(result.last_hidden_state.shape	,							(self.batch_size, self.seq_length, self.hidden_size)      )
												self.parent.assertEqual(result.pooler_output.shape	,							(self.batch_size, self.hidden_size)      )




							def 	__snake_case(						self   :					Dict	,							_UpperCamelCase   :					Any	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					List[Any]	,							_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					Optional[int]	,							_UpperCamelCase   :					Dict	,							_UpperCamelCase   :					List[Any]	,							_UpperCamelCase   :					Any	,							_UpperCamelCase   :					int	,							)					->				Optional[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       True
												SCREAMING_SNAKE_CASE    			=       NezhaModel(_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(
												    _UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							encoder_hidden_states=_UpperCamelCase	,							encoder_attention_mask=_UpperCamelCase	,							)
												SCREAMING_SNAKE_CASE    			=       model(
												    _UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							encoder_hidden_states=_UpperCamelCase	,							)
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase      )
												self.parent.assertEqual(result.last_hidden_state.shape	,							(self.batch_size, self.seq_length, self.hidden_size)      )
												self.parent.assertEqual(result.pooler_output.shape	,							(self.batch_size, self.hidden_size)      )




							def 	__snake_case(						self   :					Tuple	,							_UpperCamelCase   :					Optional[int]	,							_UpperCamelCase   :					Union[str, Any]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					int      )					->				Union[str, Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaForMaskedLM(config=_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							labels=_UpperCamelCase      )
												self.parent.assertEqual(result.logits.shape	,							(self.batch_size, self.seq_length, self.vocab_size)      )




							def 	__snake_case(						self   :					int	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					Optional[int]	,							_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					int	,							_UpperCamelCase   :					Union[str, Any]	,							_UpperCamelCase   :					Any      )					->				Any:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaForNextSentencePrediction(config=_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(
												    _UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							labels=_UpperCamelCase	,							)
												self.parent.assertEqual(result.logits.shape	,							(self.batch_size, 2)      )




							def 	__snake_case(						self   :					int	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					Union[str, Any]	,							_UpperCamelCase   :					int	,							_UpperCamelCase   :					List[Any]	,							_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					int      )					->				List[Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaForPreTraining(config=_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(
												    _UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							labels=_UpperCamelCase	,							next_sentence_label=_UpperCamelCase	,							)
												self.parent.assertEqual(result.prediction_logits.shape	,							(self.batch_size, self.seq_length, self.vocab_size)      )
												self.parent.assertEqual(result.seq_relationship_logits.shape	,							(self.batch_size, 2)      )




							def 	__snake_case(						self   :					int	,							_UpperCamelCase   :					int	,							_UpperCamelCase   :					Any	,							_UpperCamelCase   :					Dict	,							_UpperCamelCase   :					Dict	,							_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					List[str]      )					->				List[Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaForQuestionAnswering(config=_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(
												    _UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							start_positions=_UpperCamelCase	,							end_positions=_UpperCamelCase	,							)
												self.parent.assertEqual(result.start_logits.shape	,							(self.batch_size, self.seq_length)      )
												self.parent.assertEqual(result.end_logits.shape	,							(self.batch_size, self.seq_length)      )




							def 	__snake_case(						self   :					List[str]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					Union[str, Any]	,							_UpperCamelCase   :					str	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					str	,							_UpperCamelCase   :					Optional[int]      )					->				int:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.num_labels
												SCREAMING_SNAKE_CASE    			=       NezhaForSequenceClassification(_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							labels=_UpperCamelCase      )
												self.parent.assertEqual(result.logits.shape	,							(self.batch_size, self.num_labels)      )




							def 	__snake_case(						self   :					Tuple	,							_UpperCamelCase   :					Optional[int]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					str	,							_UpperCamelCase   :					str	,							_UpperCamelCase   :					str	,							_UpperCamelCase   :					Union[str, Any]	,							_UpperCamelCase   :					Tuple      )					->				List[str]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.num_labels
												SCREAMING_SNAKE_CASE    			=       NezhaForTokenClassification(config=_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							attention_mask=_UpperCamelCase	,							token_type_ids=_UpperCamelCase	,							labels=_UpperCamelCase      )
												self.parent.assertEqual(result.logits.shape	,							(self.batch_size, self.seq_length, self.num_labels)      )




							def 	__snake_case(						self   :					Union[str, Any]	,							_UpperCamelCase   :					str	,							_UpperCamelCase   :					List[Any]	,							_UpperCamelCase   :					Dict	,							_UpperCamelCase   :					Optional[Any]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					Tuple      )					->				List[str]:



												'''simple docstring'''

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





							def 	__snake_case(						self   :					Any      )					->				List[str]:



												'''simple docstring'''

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



@require_torch
class 						lowercase    (     a				,       a				,       a				,       unittest.TestCase     ):
							lowercase__						:       List[Any]								=							(
							    (
							        NezhaModel,
							        NezhaForMaskedLM,
							        NezhaForMultipleChoice,
							        NezhaForNextSentencePrediction,
							        NezhaForPreTraining,
							        NezhaForQuestionAnswering,
							        NezhaForSequenceClassification,
							        NezhaForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							lowercase__						:       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__						:       Any								=							True




							def 	__snake_case(						self   :					Optional[Any]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					List[str]	,							_UpperCamelCase   :					List[Any]=False      )					->				str:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       super()._prepare_for_class(_UpperCamelCase	,							_UpperCamelCase	,							return_labels=_UpperCamelCase      )

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




							def 	__snake_case(						self   :					Tuple      )					->				Any:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaModelTester(self      )
												SCREAMING_SNAKE_CASE    			=       ConfigTester(self	,							config_class=_UpperCamelCase	,							hidden_size=37      )




							def 	__snake_case(						self   :					List[str]      )					->				str:



												'''simple docstring'''

												self.config_tester.run_common_tests()




							def 	__snake_case(						self   :					Union[str, Any]      )					->				Tuple:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_UpperCamelCase      )




							def 	__snake_case(						self   :					Optional[Any]      )					->				Any:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs_for_decoder()
												self.model_tester.create_and_check_model_as_decoder(*_UpperCamelCase      )




							def 	__snake_case(						self   :					Dict      )					->				Any:



												'''simple docstring'''

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

												SCREAMING_SNAKE_CASE    			=       None

												self.model_tester.create_and_check_model_as_decoder(
												    _UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							)




							def 	__snake_case(						self   :					Optional[Any]      )					->				Tuple:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase      )




							def 	__snake_case(						self   :					int      )					->				str:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_multiple_choice(*_UpperCamelCase      )




							def 	__snake_case(						self   :					Tuple      )					->				int:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_next_sequence_prediction(*_UpperCamelCase      )




							def 	__snake_case(						self   :					Optional[int]      )					->				List[str]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_pretraining(*_UpperCamelCase      )




							def 	__snake_case(						self   :					str      )					->				Optional[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_question_answering(*_UpperCamelCase      )




							def 	__snake_case(						self   :					int      )					->				List[Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_sequence_classification(*_UpperCamelCase      )




							def 	__snake_case(						self   :					int      )					->				int:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase      )




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



												'''simple docstring'''

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





							@slow
							@require_torch_gpu
							def 	__snake_case(						self   :					Any      )					->				Any:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE			,   SCREAMING_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

																	SCREAMING_SNAKE_CASE    			=       True
																	SCREAMING_SNAKE_CASE    			=       model_class(config=_UpperCamelCase      )

																	SCREAMING_SNAKE_CASE    			=       self._prepare_for_class(_UpperCamelCase	,							_UpperCamelCase      )
																	SCREAMING_SNAKE_CASE    			=       torch.jit.trace(
																	    _UpperCamelCase	,							(inputs_dict["input_ids"].to("cpu"      ), inputs_dict["attention_mask"].to("cpu"      ))      )

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



@require_torch
class 						lowercase    (     unittest.TestCase     ):




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



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaModel.from_pretrained("sijunhe/nezha-cn-base"      )
												SCREAMING_SNAKE_CASE    			=       torch.tensor([[0, 1, 2, 3, 4, 5]]      )
												SCREAMING_SNAKE_CASE    			=       torch.tensor([[0, 1, 1, 1, 1, 1]]      )
												with torch.no_grad():
																	SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							attention_mask=_UpperCamelCase      )[0]
												SCREAMING_SNAKE_CASE    			=       torch.Size((1, 6, 768)      )
												self.assertEqual(output.shape	,							_UpperCamelCase      )
												SCREAMING_SNAKE_CASE    			=       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]	,							_UpperCamelCase	,							atol=1e-4      )      )





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



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       NezhaForMaskedLM.from_pretrained("sijunhe/nezha-cn-base"      )
												SCREAMING_SNAKE_CASE    			=       torch.tensor([[0, 1, 2, 3, 4, 5]]      )
												SCREAMING_SNAKE_CASE    			=       torch.tensor([[1, 1, 1, 1, 1, 1]]      )
												with torch.no_grad():
																	SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase	,							attention_mask=_UpperCamelCase      )[0]
												SCREAMING_SNAKE_CASE    			=       torch.Size((1, 6, 21_128)      )
												self.assertEqual(output.shape	,							_UpperCamelCase      )
												SCREAMING_SNAKE_CASE    			=       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]	,							_UpperCamelCase	,							atol=1e-4      )      )





import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


_lowerCamelCase					:   Union[str, Any]							=		logging.get_logger(__name__)


class 						lowercase    (     a     ):
							lowercase__						:       Optional[Any]								=							["""input_features""", """is_longer"""]




							def __init__(						self   :					str	,							_UpperCamelCase   :					Optional[int]=64	,							_UpperCamelCase   :					Any=48_000	,							_UpperCamelCase   :					Optional[Any]=480	,							_UpperCamelCase   :					List[Any]=10	,							_UpperCamelCase   :					Any=1_024	,							_UpperCamelCase   :					List[Any]=0.0	,							_UpperCamelCase   :					Any=False	,							_UpperCamelCase   :					float = 0	,							_UpperCamelCase   :					float = 14_000	,							_UpperCamelCase   :					int = None	,							_UpperCamelCase   :					str = "fusion"	,							_UpperCamelCase   :					str = "repeatpad"	,							**_UpperCamelCase   :					Optional[Any]	,							)					->				Dict:



												'''simple docstring'''

												super().__init__(
												    feature_size=_UpperCamelCase	,							sampling_rate=_UpperCamelCase	,							padding_value=_UpperCamelCase	,							return_attention_mask=_UpperCamelCase	,							**_UpperCamelCase	,							)
												SCREAMING_SNAKE_CASE    			=       top_db
												SCREAMING_SNAKE_CASE    			=       truncation
												SCREAMING_SNAKE_CASE    			=       padding
												SCREAMING_SNAKE_CASE    			=       fft_window_size
												SCREAMING_SNAKE_CASE    			=       (fft_window_size >> 1) + 1
												SCREAMING_SNAKE_CASE    			=       hop_length
												SCREAMING_SNAKE_CASE    			=       max_length_s
												SCREAMING_SNAKE_CASE    			=       max_length_s * sampling_rate
												SCREAMING_SNAKE_CASE    			=       sampling_rate
												SCREAMING_SNAKE_CASE    			=       frequency_min
												SCREAMING_SNAKE_CASE    			=       frequency_max
												SCREAMING_SNAKE_CASE    			=       mel_filter_bank(
												    num_frequency_bins=self.nb_frequency_bins	,							num_mel_filters=_UpperCamelCase	,							min_frequency=_UpperCamelCase	,							max_frequency=_UpperCamelCase	,							sampling_rate=_UpperCamelCase	,							norm=_UpperCamelCase	,							mel_scale="htk"	,							)
												SCREAMING_SNAKE_CASE    			=       mel_filter_bank(
												    num_frequency_bins=self.nb_frequency_bins	,							num_mel_filters=_UpperCamelCase	,							min_frequency=_UpperCamelCase	,							max_frequency=_UpperCamelCase	,							sampling_rate=_UpperCamelCase	,							norm="slaney"	,							mel_scale="slaney"	,							)




							def 	__snake_case(						self   :					str      )					->				Dict[str, Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       copy.deepcopy(self.__dict__      )
												SCREAMING_SNAKE_CASE    			=       self.__class__.__name__
												if "mel_filters" in output:
																	del output["mel_filters"]
												if "mel_filters_slaney" in output:
																	del output["mel_filters_slaney"]
												return output




							def 	__snake_case(						self   :					Optional[Any]	,							_UpperCamelCase   :					np.array	,							_UpperCamelCase   :					Optional[np.array] = None      )					->				np.ndarray:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       spectrogram(
												    _UpperCamelCase	,							window_function(self.fft_window_size	,							"hann"      )	,							frame_length=self.fft_window_size	,							hop_length=self.hop_length	,							power=2.0	,							mel_filters=_UpperCamelCase	,							log_mel="dB"	,							)
												return log_mel_spectrogram.T




							def 	__snake_case(						self   :					str	,							_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					Any	,							_UpperCamelCase   :					str      )					->				Tuple:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       np.array_split(list(range(0	,							total_frames - chunk_frames + 1      )      )	,							3      )
												if len(ranges[1]      ) == 0:
																	# if the audio is too short, we just use the first chunk
																	SCREAMING_SNAKE_CASE    			=       [0]
												if len(ranges[2]      ) == 0:
																	# if the audio is too short, we just use the first chunk
																	SCREAMING_SNAKE_CASE    			=       [0]
												# randomly choose index for each part
												SCREAMING_SNAKE_CASE    			=       np.random.choice(ranges[0]      )
												SCREAMING_SNAKE_CASE    			=       np.random.choice(ranges[1]      )
												SCREAMING_SNAKE_CASE    			=       np.random.choice(ranges[2]      )

												SCREAMING_SNAKE_CASE    			=       mel[idx_front : idx_front + chunk_frames, :]
												SCREAMING_SNAKE_CASE    			=       mel[idx_middle : idx_middle + chunk_frames, :]
												SCREAMING_SNAKE_CASE    			=       mel[idx_back : idx_back + chunk_frames, :]

												SCREAMING_SNAKE_CASE    			=       torch.tensor(mel[None, None, :]      )
												SCREAMING_SNAKE_CASE    			=       torch.nn.functional.interpolate(
												    _UpperCamelCase	,							size=[chunk_frames, 64]	,							mode="bilinear"	,							align_corners=_UpperCamelCase      )
												SCREAMING_SNAKE_CASE    			=       mel_shrink[0][0].numpy()
												SCREAMING_SNAKE_CASE    			=       np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]	,							axis=0      )
												return mel_fusion




							def 	__snake_case(						self   :					Optional[int]	,							_UpperCamelCase   :					np.array	,							_UpperCamelCase   :					Dict	,							_UpperCamelCase   :					Union[str, Any]	,							_UpperCamelCase   :					List[str]      )					->				np.array:



												'''simple docstring'''

												if waveform.shape[0] > max_length:
																	if truncation == "rand_trunc":
																						SCREAMING_SNAKE_CASE    			=       True
																						# random crop to max_length (for compatibility) -> this should be handled by self.pad
																						SCREAMING_SNAKE_CASE    			=       len(_UpperCamelCase      ) - max_length
																						SCREAMING_SNAKE_CASE    			=       np.random.randint(0	,							overflow + 1      )
																						SCREAMING_SNAKE_CASE    			=       waveform[idx : idx + max_length]
																						SCREAMING_SNAKE_CASE    			=       self._np_extract_fbank_features(_UpperCamelCase	,							self.mel_filters_slaney      )[None, :]
																	elif truncation == "fusion":
																						SCREAMING_SNAKE_CASE    			=       self._np_extract_fbank_features(_UpperCamelCase	,							self.mel_filters      )
																						SCREAMING_SNAKE_CASE    			=       max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
																						SCREAMING_SNAKE_CASE    			=       mel.shape[0]
																						if chunk_frames == total_frames:
																											# there is a corner case where the audio length is larger than max_length but smaller than max_length+hop_length.
																											# In this case, we just use the whole audio.
																											SCREAMING_SNAKE_CASE    			=       np.stack([mel, mel, mel, mel]	,							axis=0      )
																											SCREAMING_SNAKE_CASE    			=       False
																						else:
																											SCREAMING_SNAKE_CASE    			=       self._random_mel_fusion(_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase      )
																											SCREAMING_SNAKE_CASE    			=       True
																	else:
																						raise NotImplementedError(F"data_truncating {truncation} not implemented"      )

												else:
																	SCREAMING_SNAKE_CASE    			=       False
																	# only use repeat as a new possible value for padding. you repeat the audio before applying the usual max_length padding
																	if waveform.shape[0] < max_length:
																						if padding == "repeat":
																											SCREAMING_SNAKE_CASE    			=       int(max_length / len(_UpperCamelCase      )      )
																											SCREAMING_SNAKE_CASE    			=       np.stack(np.tile(_UpperCamelCase	,							n_repeat + 1      )      )[:max_length]
																						if padding == "repeatpad":
																											SCREAMING_SNAKE_CASE    			=       int(max_length / len(_UpperCamelCase      )      )
																											SCREAMING_SNAKE_CASE    			=       np.stack(np.tile(_UpperCamelCase	,							_UpperCamelCase      )      )
																						SCREAMING_SNAKE_CASE    			=       np.pad(_UpperCamelCase	,							(0, max_length - waveform.shape[0])	,							mode="constant"	,							constant_values=0      )

																	if truncation == "fusion":
																						SCREAMING_SNAKE_CASE    			=       self._np_extract_fbank_features(_UpperCamelCase	,							self.mel_filters      )
																						SCREAMING_SNAKE_CASE    			=       np.stack([input_mel, input_mel, input_mel, input_mel]	,							axis=0      )
																	else:
																						SCREAMING_SNAKE_CASE    			=       self._np_extract_fbank_features(_UpperCamelCase	,							self.mel_filters_slaney      )[None, :]

												return input_mel, longer





							def __call__(						self   :					Dict	,							_UpperCamelCase   :					Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]	,							_UpperCamelCase   :					str = None	,							_UpperCamelCase   :					Optional[str] = None	,							_UpperCamelCase   :					Optional[int] = None	,							_UpperCamelCase   :					Optional[int] = None	,							_UpperCamelCase   :					Optional[Union[str, TensorType]] = None	,							**_UpperCamelCase   :					Tuple	,							)					->				BatchFeature:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       truncation if truncation is not None else self.truncation
												SCREAMING_SNAKE_CASE    			=       padding if padding else self.padding

												if sampling_rate is not None:
																	if sampling_rate != self.sampling_rate:
																						raise ValueError(
																						    F"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"
																						    F" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"
																						    F" was sampled with {self.sampling_rate} and not {sampling_rate}."      )
												else:
																	logger.warning(
																	    "It is strongly recommended to pass the `sampling_rate` argument to this function. "
																	    "Failing to do so can result in silent errors that might be hard to debug."      )

												SCREAMING_SNAKE_CASE    			=       isinstance(_UpperCamelCase	,							np.ndarray      ) and len(raw_speech.shape      ) > 1
												if is_batched_numpy and len(raw_speech.shape      ) > 2:
																	raise ValueError(F"Only mono-channel audio is supported for input to {self}"      )
												SCREAMING_SNAKE_CASE    			=       is_batched_numpy or (
												    isinstance(_UpperCamelCase	,							(list, tuple)      ) and (isinstance(raw_speech[0]	,							(np.ndarray, tuple, list)      ))
												)

												if is_batched:
																	SCREAMING_SNAKE_CASE    			=       [np.asarray(_UpperCamelCase	,							dtype=np.floataa      ) for speech in raw_speech]
												elif not is_batched and not isinstance(_UpperCamelCase	,							np.ndarray      ):
																	SCREAMING_SNAKE_CASE    			=       np.asarray(_UpperCamelCase	,							dtype=np.floataa      )
												elif isinstance(_UpperCamelCase	,							np.ndarray      ) and raw_speech.dtype is np.dtype(np.floataa      ):
																	SCREAMING_SNAKE_CASE    			=       raw_speech.astype(np.floataa      )

												# always return batch
												if not is_batched:
																	SCREAMING_SNAKE_CASE    			=       [np.asarray(_UpperCamelCase      )]

												# convert to mel spectrogram, truncate and pad if needed.
												SCREAMING_SNAKE_CASE    			=       [
												    self._get_input_mel(_UpperCamelCase	,							max_length if max_length else self.nb_max_samples	,							_UpperCamelCase	,							_UpperCamelCase      )
												    for waveform in raw_speech
												]

												SCREAMING_SNAKE_CASE    			=       []
												SCREAMING_SNAKE_CASE    			=       []
												for mel, longer in padded_inputs:
																	input_mel.append(_UpperCamelCase      )
																	is_longer.append(_UpperCamelCase      )

												if truncation == "fusion" and sum(_UpperCamelCase      ) == 0:
																	# if no audio is longer than 10s, then randomly select one audio to be longer
																	SCREAMING_SNAKE_CASE    			=       np.random.randint(0	,							len(_UpperCamelCase      )      )
																	SCREAMING_SNAKE_CASE    			=       True

												if isinstance(input_mel[0]	,							_UpperCamelCase      ):
																	SCREAMING_SNAKE_CASE    			=       [np.asarray(_UpperCamelCase	,							dtype=np.floataa      ) for feature in input_mel]

												# is_longer is a list of bool
												SCREAMING_SNAKE_CASE    			=       [[longer] for longer in is_longer]

												SCREAMING_SNAKE_CASE    			=       {"input_features": input_mel, "is_longer": is_longer}
												SCREAMING_SNAKE_CASE    			=       BatchFeature(_UpperCamelCase      )

												if return_tensors is not None:
																	SCREAMING_SNAKE_CASE    			=       input_features.convert_to_tensors(_UpperCamelCase      )

												return input_features




