


from __future__ import annotations

import unittest

from transformers import OpenAIGPTConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers.models.openai.modeling_tf_openai import (
           TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFOpenAIGPTDoubleHeadsModel,
           TFOpenAIGPTForSequenceClassification,
           TFOpenAIGPTLMHeadModel,
           TFOpenAIGPTModel,
       )


class 	a__						:

       def __init__( self   :  Optional[int]   ,				lowerCamelCase_   :  Optional[Any]   ,				):
            a_     :  Any				      =     parent
            a_     :  int				      =     1_3
            a_     :  Any				      =     7
            a_     :  Any				      =     True
            a_     :  Tuple				      =     True
            a_     :  Optional[Any]				      =     True
            a_     :  List[Any]				      =     True
            a_     :  Any				      =     True
            a_     :  str				      =     9_9
            a_     :  Tuple				      =     3_2
            a_     :  Union[str, Any]				      =     2
            a_     :  Union[str, Any]				      =     4
            a_     :  Optional[Any]				      =     3_7
            a_     :  Optional[int]				      =     """gelu"""
            a_     :  Union[str, Any]				      =     0.1
            a_     :  int				      =     0.1
            a_     :  Union[str, Any]				      =     5_1_2
            a_     :  List[Any]				      =     1_6
            a_     :  Any				      =     2
            a_     :  Dict				      =     0.0_2
            a_     :  Optional[int]				      =     3
            a_     :  Optional[int]				      =     4
            a_     :  str				      =     None
            a_     :  Dict				      =     self.vocab_size - 1

       def     UpperCAmelCase( self   :  List[Any]     ):
            a_     :  Union[str, Any]				      =     ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size     )

            a_     :  int				      =     None
            if self.use_input_mask:
                 a_     :  List[str]				      =     random_attention_mask([self.batch_size, self.seq_length]     )

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

            a_     :  List[Any]				      =     None
            if self.use_mc_token_ids:
                 a_     :  Dict				      =     ids_tensor([self.batch_size, self.num_choices]   ,				self.seq_length     )

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

            a_     :  Any				      =     OpenAIGPTConfig(
                vocab_size=self.vocab_size   ,				n_embd=self.hidden_size   ,				n_layer=self.num_hidden_layers   ,				n_head=self.num_attention_heads   ,				n_positions=self.max_position_embeddings   ,				pad_token_id=self.pad_token_id   ,				)

            a_     :  int				      =     ids_tensor([self.num_hidden_layers, self.num_attention_heads]   ,				2     )

            return (
                config,
                input_ids,
                input_mask,
                head_mask,
                token_type_ids,
                mc_token_ids,
                sequence_labels,
                token_labels,
                choice_labels,
            )

       def     UpperCAmelCase( self   :  List[str]   ,				lowerCamelCase_   :  Tuple   ,				lowerCamelCase_   :  List[Any]   ,				lowerCamelCase_   :  str   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  List[Any]   ,				*lowerCamelCase_   :  Any     ):
            a_     :  Any				      =     TFOpenAIGPTModel(config=__UpperCamelCase     )
            a_     :  List[str]				      =     {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
            a_     :  int				      =     model(__UpperCamelCase     )

            a_     :  Optional[Any]				      =     [input_ids, input_mask]
            a_     :  int				      =     model(__UpperCamelCase     )

            a_     :  List[Any]				      =     model(__UpperCamelCase     )

            self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.hidden_size)     )

       def     UpperCAmelCase( self   :  Optional[Any]   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  str   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  int   ,				*lowerCamelCase_   :  Union[str, Any]     ):
            a_     :  Optional[int]				      =     TFOpenAIGPTLMHeadModel(config=__UpperCamelCase     )
            a_     :  str				      =     {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
            a_     :  Optional[int]				      =     model(__UpperCamelCase     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.vocab_size)     )

       def     UpperCAmelCase( self   :  List[Any]   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  str   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  str   ,				*lowerCamelCase_   :  str     ):
            a_     :  Tuple				      =     TFOpenAIGPTDoubleHeadsModel(config=__UpperCamelCase     )

            a_     :  Optional[int]				      =     tf.tile(tf.expand_dims(__UpperCamelCase   ,				1     )   ,				(1, self.num_choices, 1)     )
            a_     :  Any				      =     tf.tile(tf.expand_dims(__UpperCamelCase   ,				1     )   ,				(1, self.num_choices, 1)     )
            a_     :  Union[str, Any]				      =     tf.tile(tf.expand_dims(__UpperCamelCase   ,				1     )   ,				(1, self.num_choices, 1)     )

            a_     :  Optional[int]				      =     {
                """input_ids""": multiple_choice_inputs_ids,
                """mc_token_ids""": mc_token_ids,
                """attention_mask""": multiple_choice_input_mask,
                """token_type_ids""": multiple_choice_token_type_ids,
            }
            a_     :  Union[str, Any]				      =     model(__UpperCamelCase     )
            self.parent.assertEqual(
                result.logits.shape   ,				(self.batch_size, self.num_choices, self.seq_length, self.vocab_size)     )
            self.parent.assertEqual(result.mc_logits.shape   ,				(self.batch_size, self.num_choices)     )

       def     UpperCAmelCase( self   :  Tuple   ,				lowerCamelCase_   :  List[Any]   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  str   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  List[str]   ,				*lowerCamelCase_   :  Tuple     ):
            a_     :  Union[str, Any]				      =     self.num_labels
            a_     :  Optional[Any]				      =     ids_tensor([self.batch_size]   ,				self.type_sequence_label_size     )
            a_     :  int				      =     {
                """input_ids""": input_ids,
                """attention_mask""": input_mask,
                """token_type_ids""": token_type_ids,
                """labels""": sequence_labels,
            }
            a_     :  Dict				      =     TFOpenAIGPTForSequenceClassification(__UpperCamelCase     )
            a_     :  Union[str, Any]				      =     model(__UpperCamelCase     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)     )

       def     UpperCAmelCase( self   :  Optional[int]     ):
            a_     :  Dict				      =     self.prepare_config_and_inputs()

            (
                (
                a_			
            )      ,      (
                a_			
            )      ,      (
                a_			
            )      ,      (
                a_			
            )      ,      (
                a_			
            )      ,      (
                a_			
            )      ,      (
                a_			
            )      ,      (
                a_			
            )      ,      (
                a_			
            )      ,      			
            )     :  Union[str, Any]				      =     config_and_inputs

            a_     :  List[str]				      =     {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
            return config, inputs_dict


@require_tf
class 	a__						(    __lowerCAmelCase    ,							__lowerCAmelCase    ,							unittest.TestCase     ):
       lowerCamelCase__:  List[Any]				     =  (
           (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
           if is_tf_available()
           else ()
       )
       lowerCamelCase__:  Optional[Any]				     =  (
           (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
       )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
       lowerCamelCase__:  Union[str, Any]				     =  (
           {
               """feature-extraction""": TFOpenAIGPTModel,
               """text-classification""": TFOpenAIGPTForSequenceClassification,
               """text-generation""": TFOpenAIGPTLMHeadModel,
               """zero-shot""": TFOpenAIGPTForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       lowerCamelCase__:  Optional[Any]				     =  False
       lowerCamelCase__:  int				     =  False

       def     UpperCAmelCase( self   :  Optional[int]   ,				lowerCamelCase_   :  str   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  Union[str, Any]     ):






            if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
                 # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
                 # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
                 # tiny config could not be created.
                 return True

            return False

       def     UpperCAmelCase( self   :  int     ):
            a_     :  str				      =     TFOpenAIGPTModelTester(self     )
            a_     :  List[str]				      =     ConfigTester(self   ,				config_class=__UpperCamelCase   ,				n_embd=3_7     )

       def     UpperCAmelCase( self   :  Any     ):
            self.config_tester.run_common_tests()

       def     UpperCAmelCase( self   :  Optional[int]     ):
            a_     :  Dict				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_model(*__UpperCamelCase     )

       def     UpperCAmelCase( self   :  Union[str, Any]     ):
            a_     :  List[Any]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_lm_head(*__UpperCamelCase     )

       def     UpperCAmelCase( self   :  int     ):
            a_     :  Dict				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_double_head(*__UpperCamelCase     )

       def     UpperCAmelCase( self   :  Tuple     ):
            a_     :  List[Any]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*__UpperCamelCase     )

       @slow
       def     UpperCAmelCase( self   :  List[Any]     ):






            for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 a_     :  Tuple				      =     TFOpenAIGPTModel.from_pretrained(__UpperCamelCase     )
                 self.assertIsNotNone(__UpperCamelCase     )


@require_tf
class 	a__						(    unittest.TestCase     ):

       @slow
       def     UpperCAmelCase( self   :  Dict     ):
            a_     :  Union[str, Any]				      =     TFOpenAIGPTLMHeadModel.from_pretrained("""openai-gpt"""     )
            a_     :  List[str]				      =     tf.convert_to_tensor([[4_8_1, 4_7_3_5, 5_4_4]]   ,				dtype=tf.intaa     )  # the president is
            a_     :  Any				      =     [
                4_8_1,
                4_7_3_5,
                5_4_4,
                2_4_6,
                9_6_3,
                8_7_0,
                7_6_2,
                2_3_9,
                2_4_4,
                4_0_4_7_7,
                2_4_4,
                2_4_9,
                7_1_9,
                8_8_1,
                4_8_7,
                5_4_4,
                2_4_0,
                2_4_4,
                6_0_3,
                4_8_1,
            ]  # the president is a very good man. " \n " i\'m sure he is, " said the

            a_     :  int				      =     model.generate(__UpperCamelCase   ,				do_sample=__UpperCamelCase     )
            self.assertListEqual(output_ids[0].numpy().tolist()   ,				__UpperCamelCase     )









'''simple docstring'''

import gc
import unittest

import numpy as np
import torch

from diffusers import DDPMScheduler, MidiProcessor, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder
from diffusers.utils import require_torch_gpu, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_note_seq, require_onnxruntime

from ..pipeline_params import TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS, TOKENS_TO_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


a__         =     '''./tests/fixtures/elise_format0.mid'''





@unittest.skip("""The note-seq package currently throws an error on import"""      )
class __magic_name__(						__lowerCAmelCase    ,		unittest.TestCase      ):
						UpperCAmelCase_		:					Tuple    =  SpectrogramDiffusionPipeline
						UpperCAmelCase_		:					List[str]    =  PipelineTesterMixin.required_optional_params - {
						    """callback""",
						    """latents""",
						    """callback_steps""",
						    """output_type""",
						    """num_images_per_prompt""",
						}
						UpperCAmelCase_		:					Dict    =  False
						UpperCAmelCase_		:					Dict    =  False
						UpperCAmelCase_		:					List[str]    =  TOKENS_TO_AUDIO_GENERATION_PARAMS
						UpperCAmelCase_		:					Union[str, Any]    =  TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS






						def       __lowerCAmelCase(   self						:						Tuple       ):



									'''simple docstring'''






									torch.manual_seed(0       )
									snake_case__          =						SpectrogramNotesEncoder(
									    max_length=2_0_4_8					,   vocab_size=1_5_3_6					,   d_model=7_6_8					,   dropout_rate=0.1					,   num_layers=1					,   num_heads=1					,   d_kv=4					,   d_ff=2_0_4_8					,   feed_forward_proj="""gated-gelu"""					,   )

									snake_case__          =						SpectrogramContEncoder(
									    input_dims=1_2_8					,   targets_context_length=2_5_6					,   d_model=7_6_8					,   dropout_rate=0.1					,   num_layers=1					,   num_heads=1					,   d_kv=4					,   d_ff=2_0_4_8					,   feed_forward_proj="""gated-gelu"""					,   )

									snake_case__          =						TaFilmDecoder(
									    input_dims=1_2_8					,   targets_length=2_5_6					,   max_decoder_noise_time=2_00_00.0					,   d_model=7_6_8					,   num_layers=1					,   num_heads=1					,   d_kv=4					,   d_ff=2_0_4_8					,   dropout_rate=0.1					,   )

									snake_case__          =						DDPMScheduler()

									snake_case__          =						{
									    """notes_encoder""": notes_encoder.eval(),
									    """continuous_encoder""": continuous_encoder.eval(),
									    """decoder""": decoder.eval(),
									    """scheduler""": scheduler,
									    """melgan""": None,
									}
									return components






						def       __lowerCAmelCase(   self						:						Tuple					,   __UpperCamelCase						:						Any					,   __UpperCamelCase						:						str=0       ):



									'''simple docstring'''






									if str(__UpperCamelCase       ).startswith("""mps"""       ):
												snake_case__          =						torch.manual_seed(__UpperCamelCase       )
									else:
												snake_case__          =						torch.Generator(device=__UpperCamelCase       ).manual_seed(__UpperCamelCase       )
									snake_case__          =						{
									    """input_tokens""": [
									        [1_1_3_4, 9_0, 1_1_3_5, 1_1_3_3, 1_0_8_0, 1_1_2, 1_1_3_2, 1_0_8_0, 1_1_3_3, 1_0_7_9, 1_3_3, 1_1_3_2, 1_0_7_9, 1_1_3_3, 1] + [0] * 2_0_3_3
									    ],
									    """generator""": generator,
									    """num_inference_steps""": 4,
									    """output_type""": """mel""",
									}
									return inputs






						def       __lowerCAmelCase(   self						:						Optional[Any]       ):



									'''simple docstring'''






									snake_case__          =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
									snake_case__          =						self.get_dummy_components()
									snake_case__          =						SpectrogramDiffusionPipeline(**__UpperCamelCase       )
									snake_case__          =						pipe.to(__UpperCamelCase       )
									pipe.set_progress_bar_config(disable=__UpperCamelCase       )

									snake_case__          =						self.get_dummy_inputs(__UpperCamelCase       )
									snake_case__          =						pipe(**__UpperCamelCase       )
									snake_case__          =						output.audios

									snake_case__          =						mel[0, -3:, -3:]

									assert mel_slice.shape == (3, 3)
									snake_case__          =						np.array(
									    [-11.51_29_25, -4.78_82_15, -0.46_17_29_05, -2.05_17_15, -10.53_91_47, -10.97_09_63, -9.09_16_34, 4.0, 4.0]       )
									assert np.abs(mel_slice.flatten() - expected_slice       ).max() < 1E-2






						@skip_mps
						def       __lowerCAmelCase(   self						:						Tuple       ):



									'''simple docstring'''






									return super().test_save_load_local()






						@skip_mps
						def       __lowerCAmelCase(   self						:						List[Any]       ):



									'''simple docstring'''






									return super().test_dict_tuple_outputs_equivalent()






						@skip_mps
						def       __lowerCAmelCase(   self						:						Optional[int]       ):



									'''simple docstring'''






									return super().test_save_load_optional_components()






						@skip_mps
						def       __lowerCAmelCase(   self						:						Union[str, Any]       ):



									'''simple docstring'''






									return super().test_attention_slicing_forward_pass()






						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									pass






						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									pass







						@skip_mps
						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									return super().test_progress_bar()







@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class __magic_name__(						unittest.TestCase      ):






						def       __lowerCAmelCase(   self						:						List[Any]       ):



									'''simple docstring'''






									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									snake_case__          =						torch_device

									snake_case__          =						SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""       )
									snake_case__          =						pipe.melgan
									snake_case__          =						None

									snake_case__          =						pipe.to(__UpperCamelCase       )
									pipe.set_progress_bar_config(disable=__UpperCamelCase       )

									def callback(__UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						List[str]       ):
												# decode mel to audio
												snake_case__          =						melgan(input_features=mel_output.astype(np.floataa       )       )[0]
												assert len(audio[0]       ) == 8_1_9_2_0 * (step + 1)
												# simulate that audio is played
												return audio

									snake_case__          =						MidiProcessor()
									snake_case__          =						processor(__UpperCamelCase       )

									snake_case__          =						input_tokens[:3]
									snake_case__          =						torch.manual_seed(0       )
									pipe(__UpperCamelCase					,   num_inference_steps=5					,   generator=__UpperCamelCase					,   callback=__UpperCamelCase					,   output_type="""mel"""       )






						def       __lowerCAmelCase(   self						:						List[Any]       ):



									'''simple docstring'''






									snake_case__          =						torch_device

									snake_case__          =						SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""       )
									snake_case__          =						pipe.to(__UpperCamelCase       )
									pipe.set_progress_bar_config(disable=__UpperCamelCase       )
									snake_case__          =						MidiProcessor()

									snake_case__          =						processor(__UpperCamelCase       )
									# just run two denoising loops
									snake_case__          =						input_tokens[:2]

									snake_case__          =						torch.manual_seed(0       )
									snake_case__          =						pipe(__UpperCamelCase					,   num_inference_steps=2					,   generator=__UpperCamelCase       )

									snake_case__          =						output.audios[0]

									assert abs(np.abs(__UpperCamelCase       ).sum() - 36_12.8_41       ) < 1E-1







						def       __lowerCAmelCase(   self						:						Union[str, Any]       ):



									'''simple docstring'''






									snake_case__          =						torch_device

									snake_case__          =						SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""       )
									snake_case__          =						pipe.to(__UpperCamelCase       )
									pipe.set_progress_bar_config(disable=__UpperCamelCase       )

									snake_case__          =						MidiProcessor()

									snake_case__          =						processor(__UpperCamelCase       )

									# just run 4 denoising loops
									snake_case__          =						input_tokens[:4]

									snake_case__          =						torch.manual_seed(0       )
									snake_case__          =						pipe(__UpperCamelCase					,   num_inference_steps=1_0_0					,   generator=__UpperCamelCase       )

									snake_case__          =						output.audios[0]
									assert abs(np.abs(__UpperCamelCase       ).sum() - 93_89.11_11       ) < 5E-2