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			,			__lowercase			,			__lowercase=13			,			__lowercase=7			,			__lowercase=True			,			__lowercase=True			,			__lowercase=True			,			__lowercase=True			,			__lowercase=99			,			__lowercase=32			,			__lowercase=5			,			__lowercase=4			,			__lowercase=37			,			__lowercase="gelu"			,			__lowercase=0.1			,			__lowercase=0.1			,			__lowercase=128			,			__lowercase=32			,			__lowercase=16			,			__lowercase=2			,			__lowercase=0.02			,			__lowercase=3			,			__lowercase=4			,			__lowercase=None			,			)	->     Tuple:
        __UpperCamelCase					:Optional[int]        =		parent
        __UpperCamelCase					:int        =		batch_size
        __UpperCamelCase					:Dict        =		seq_length
        __UpperCamelCase					:Any        =		is_training
        __UpperCamelCase					:str        =		use_input_mask
        __UpperCamelCase					:Dict        =		use_token_type_ids
        __UpperCamelCase					:Union[str, Any]        =		use_labels
        __UpperCamelCase					:Tuple        =		vocab_size
        __UpperCamelCase					:int        =		hidden_size
        __UpperCamelCase					:Dict        =		num_hidden_layers
        __UpperCamelCase					:Optional[int]        =		num_attention_heads
        __UpperCamelCase					:List[str]        =		intermediate_size
        __UpperCamelCase					:Optional[Any]        =		hidden_act
        __UpperCamelCase					:Any        =		hidden_dropout_prob
        __UpperCamelCase					:List[str]        =		attention_probs_dropout_prob
        __UpperCamelCase					:Optional[int]        =		max_position_embeddings
        __UpperCamelCase					:Union[str, Any]        =		type_vocab_size
        __UpperCamelCase					:List[Any]        =		type_sequence_label_size
        __UpperCamelCase					:int        =		initializer_range
        __UpperCamelCase					:List[str]        =		num_labels
        __UpperCamelCase					:Optional[int]        =		num_choices
        __UpperCamelCase					:Tuple        =		scope







    def       UpperCamelCase__       (     self)	->     Any:
        __UpperCamelCase					:str        =		ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size)

        __UpperCamelCase					:Optional[Any]        =		None
        if self.use_input_mask:
            __UpperCamelCase					:List[str]        =		random_attention_mask([self.batch_size, self.seq_length])

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

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

        __UpperCamelCase					:Dict        =		self.get_config()

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







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







    def       UpperCamelCase__       (     self)	->     Union[str, Any]:
        (
            (
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	
        )					:Tuple        =		self.prepare_config_and_inputs()

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







    def       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)	->     List[Any]:
        __UpperCamelCase					:int        =		NezhaForNextSentencePrediction(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase					:List[str]        =		model(
            __lowercase			,			attention_mask=__lowercase			,			token_type_ids=__lowercase			,			labels=__lowercase			,			)
        self.parent.assertEqual(result.logits.shape			,			(self.batch_size, 2))







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







    def       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)	->     List[str]:
        __UpperCamelCase					:List[str]        =		self.num_labels
        __UpperCamelCase					:Union[str, Any]        =		NezhaForTokenClassification(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase					:Union[str, Any]        =		model(__lowercase			,			attention_mask=__lowercase			,			token_type_ids=__lowercase			,			labels=__lowercase)
        self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.num_labels))







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







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


@require_torch
class      lowerCamelCase_					(       UpperCAmelCase_						,     UpperCAmelCase_						,     UpperCAmelCase_						,     unittest.TestCase			):






    '''simple docstring'''






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







    def       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase			,			__lowercase=False)	->     int:
        __UpperCamelCase					:Union[str, Any]        =		super()._prepare_for_class(__lowercase			,			__lowercase			,			return_labels=__lowercase)

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







    def       UpperCamelCase__       (     self)	->     int:
        __UpperCamelCase					:Tuple        =		NezhaModelTester(self)
        __UpperCamelCase					:List[str]        =		ConfigTester(self			,			config_class=__lowercase			,			hidden_size=37)







    def       UpperCamelCase__       (     self)	->     Tuple:
        self.config_tester.run_common_tests()







    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:Union[str, Any]        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase)







    def       UpperCamelCase__       (     self)	->     Optional[int]:
        __UpperCamelCase					:Any        =		self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*__lowercase)







    def       UpperCamelCase__       (     self)	->     Dict:
        # This regression test was failing with PyTorch < 1.3
        (
            (
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	
        )					:List[str]        =		self.model_tester.prepare_config_and_inputs_for_decoder()

        __UpperCamelCase					:Dict        =		None

        self.model_tester.create_and_check_model_as_decoder(
            __lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			)







    def       UpperCamelCase__       (     self)	->     List[Any]:
        __UpperCamelCase					:Dict        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowercase)







    def       UpperCamelCase__       (     self)	->     str:
        __UpperCamelCase					:Dict        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowercase)







    def       UpperCamelCase__       (     self)	->     int:
        __UpperCamelCase					:Optional[int]        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_next_sequence_prediction(*__lowercase)







    def       UpperCamelCase__       (     self)	->     Optional[Any]:
        __UpperCamelCase					:Optional[int]        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*__lowercase)







    def       UpperCamelCase__       (     self)	->     Optional[int]:
        __UpperCamelCase					:Optional[int]        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowercase)







    def       UpperCamelCase__       (     self)	->     Tuple:
        __UpperCamelCase					:Union[str, Any]        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowercase)







    def       UpperCamelCase__       (     self)	->     Any:
        __UpperCamelCase					:Any        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowercase)







    @slow
    def       UpperCamelCase__       (     self)	->     str:
        for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCamelCase					:List[str]        =		NezhaModel.from_pretrained(__lowercase)
            self.assertIsNotNone(__lowercase)







    @slow
    @require_torch_gpu
    def       UpperCamelCase__       (     self)	->     str:
        __UpperCamelCase  ,	__UpperCamelCase					:Union[str, 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

            __UpperCamelCase					:str        =		True
            __UpperCamelCase					:Optional[Any]        =		model_class(config=__lowercase)

            __UpperCamelCase					:Tuple        =		self._prepare_for_class(__lowercase			,			__lowercase)
            __UpperCamelCase					:Any        =		torch.jit.trace(
                __lowercase			,			(inputs_dict['''input_ids'''].to('''cpu'''), inputs_dict['''attention_mask'''].to('''cpu''')))

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


@require_torch
class      lowerCamelCase_					(       unittest.TestCase			):






    '''simple docstring'''






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

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







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

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






from ...utils import is_note_seq_available, is_transformers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


try:
    if not (is_transformers_available() and is_torch_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
    from .notes_encoder import SpectrogramNotesEncoder
    from .continous_encoder import SpectrogramContEncoder
    from .pipeline_spectrogram_diffusion import (
        SpectrogramContEncoder,
        SpectrogramDiffusionPipeline,
        TaFilmDecoder,
    )

try:
    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
    from .midi_utils import MidiProcessor





