



from copy import deepcopy
from unittest.case import TestCase

import pytest

from datasets.arrow_dataset import Dataset
from datasets.features import Audio, ClassLabel, Features, Image, Sequence, Value
from datasets.info import DatasetInfo
from datasets.tasks import (
    AudioClassification,
    AutomaticSpeechRecognition,
    ImageClassification,
    LanguageModeling,
    QuestionAnsweringExtractive,
    Summarization,
    TextClassification,
    task_template_from_dict,
)
from datasets.utils.py_utils import asdict


A__       =			{
    """id""": """5733be284776f41900661182""",
    """title""": """University_of_Notre_Dame""",
    """context""": """Architecturally, the school has a Catholic character. Atop the Main Building\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend \"Venite Ad Me Omnes\". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.""",
    """question""": """To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?""",
    """answers""": {"""text""": ["""Saint Bernadette Soubirous"""], """answer_start""": [5_15]},
}


@pytest.mark.parametrize(
    """task_cls"""      ,		[
        AudioClassification,
        AutomaticSpeechRecognition,
        ImageClassification,
        LanguageModeling,
        QuestionAnsweringExtractive,
        Summarization,
        TextClassification,
    ]      ,		)
def 						_UpperCAmelCase							( snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       task_cls()
					_lowerCAmelCase         =       asdict(snake_case     )
					_lowerCAmelCase         =       task_template_from_dict(snake_case     )
					assert task == reloaded





class    __lowerCAmelCase					:




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       LanguageModeling(text_column="""input_text"""						)
									assert {"input_text": "text"} == task.column_mapping




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""text""": Value("""string"""						)}						)
									_lowerCAmelCase         =       {"""text_column""": """input_text"""}
									_lowerCAmelCase         =       LanguageModeling.from_dict(_snake_case						)
									assert "language-modeling" == task.task
									assert input_schema == task.input_schema





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       sorted(["""pos""", """neg"""]						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       TextClassification(text_column="""input_text"""		,							label_column="""input_label"""						)
									self.assertDictEqual({"""input_text""": """text""", """input_label""": """labels"""}		,							task.column_mapping						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""text""": Value("""string"""						)}						)
									# Labels are cast to tuple during `TextClassification.__post_init__`, so we do the same here
									_lowerCAmelCase         =       Features({"""labels""": ClassLabel}						)
									_lowerCAmelCase         =       {"""text_column""": """input_text""", """label_column""": """input_labels"""}
									_lowerCAmelCase         =       TextClassification.from_dict(_snake_case						)
									self.assertEqual("""text-classification"""		,							task.task						)
									self.assertEqual(_snake_case		,							task.input_schema						)
									self.assertEqual(_snake_case		,							task.label_schema						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       TextClassification(text_column="""input_text"""		,							label_column="""input_label"""						)
									self.assertEqual(task.label_schema["""labels"""]		,							_snake_case						)
									_lowerCAmelCase         =       task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels						)}						)						)
									self.assertEqual(task.label_schema["""labels"""]		,							ClassLabel(names=self.labels						)						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       QuestionAnsweringExtractive(
									    context_column="""input_context"""		,							question_column="""input_question"""		,							answers_column="""input_answers"""						)
									self.assertDictEqual(
									    {"""input_context""": """context""", """input_question""": """question""", """input_answers""": """answers"""}		,							task.column_mapping						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""question""": Value("""string"""						), """context""": Value("""string"""						)}						)
									_lowerCAmelCase         =       Features(
									    {
									        """answers""": Sequence(
									            {
									                """text""": Value("""string"""						),
									                """answer_start""": Value("""int32"""						),
									            }						)
									    }						)
									_lowerCAmelCase         =       {
									    """context_column""": """input_input_context""",
									    """question_column""": """input_question""",
									    """answers_column""": """input_answers""",
									}
									_lowerCAmelCase         =       QuestionAnsweringExtractive.from_dict(_snake_case						)
									self.assertEqual("""question-answering-extractive"""		,							task.task						)
									self.assertEqual(_snake_case		,							task.input_schema						)
									self.assertEqual(_snake_case		,							task.label_schema						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Summarization(text_column="""input_text"""		,							summary_column="""input_summary"""						)
									self.assertDictEqual({"""input_text""": """text""", """input_summary""": """summary"""}		,							task.column_mapping						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""text""": Value("""string"""						)}						)
									_lowerCAmelCase         =       Features({"""summary""": Value("""string"""						)}						)
									_lowerCAmelCase         =       {"""text_column""": """input_text""", """summary_column""": """input_summary"""}
									_lowerCAmelCase         =       Summarization.from_dict(_snake_case						)
									self.assertEqual("""summarization"""		,							task.task						)
									self.assertEqual(_snake_case		,							task.input_schema						)
									self.assertEqual(_snake_case		,							task.label_schema						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       AutomaticSpeechRecognition(audio_column="""input_audio"""		,							transcription_column="""input_transcription"""						)
									self.assertDictEqual({"""input_audio""": """audio""", """input_transcription""": """transcription"""}		,							task.column_mapping						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""audio""": Audio()}						)
									_lowerCAmelCase         =       Features({"""transcription""": Value("""string"""						)}						)
									_lowerCAmelCase         =       {
									    """audio_column""": """input_audio""",
									    """transcription_column""": """input_transcription""",
									}
									_lowerCAmelCase         =       AutomaticSpeechRecognition.from_dict(_snake_case						)
									self.assertEqual("""automatic-speech-recognition"""		,							task.task						)
									self.assertEqual(_snake_case		,							task.input_schema						)
									self.assertEqual(_snake_case		,							task.label_schema						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       sorted(["""pos""", """neg"""]						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       AudioClassification(audio_column="""input_audio"""		,							label_column="""input_label"""						)
									self.assertDictEqual({"""input_audio""": """audio""", """input_label""": """labels"""}		,							task.column_mapping						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""audio""": Audio()}						)
									_lowerCAmelCase         =       Features({"""labels""": ClassLabel}						)
									_lowerCAmelCase         =       {
									    """audio_column""": """input_image""",
									    """label_column""": """input_label""",
									}
									_lowerCAmelCase         =       AudioClassification.from_dict(_snake_case						)
									self.assertEqual("""audio-classification"""		,							task.task						)
									self.assertEqual(_snake_case		,							task.input_schema						)
									self.assertEqual(_snake_case		,							task.label_schema						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       AudioClassification(audio_column="""input_audio"""		,							label_column="""input_label"""						)
									self.assertEqual(task.label_schema["""labels"""]		,							_snake_case						)
									_lowerCAmelCase         =       task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels						)}						)						)
									self.assertEqual(task.label_schema["""labels"""]		,							ClassLabel(names=self.labels						)						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       sorted(["""pos""", """neg"""]						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       ImageClassification(image_column="""input_image"""		,							label_column="""input_label"""						)
									self.assertDictEqual({"""input_image""": """image""", """input_label""": """labels"""}		,							task.column_mapping						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""image""": Image()}						)
									_lowerCAmelCase         =       Features({"""labels""": ClassLabel}						)
									_lowerCAmelCase         =       {
									    """image_column""": """input_image""",
									    """label_column""": """input_label""",
									}
									_lowerCAmelCase         =       ImageClassification.from_dict(_snake_case						)
									self.assertEqual("""image-classification"""		,							task.task						)
									self.assertEqual(_snake_case		,							task.input_schema						)
									self.assertEqual(_snake_case		,							task.label_schema						)




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       ImageClassification(image_column="""input_image"""		,							label_column="""input_label"""						)
									self.assertEqual(task.label_schema["""labels"""]		,							_snake_case						)
									_lowerCAmelCase         =       task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels						)}						)						)
									self.assertEqual(task.label_schema["""labels"""]		,							ClassLabel(names=self.labels						)						)





class    __lowerCAmelCase					(	lowerCamelCase__			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       DatasetInfo(task_templates=QuestionAnsweringExtractive()						)
									_lowerCAmelCase         =       Dataset.from_dict({k: [v] for k, v in SAMPLE_QUESTION_ANSWERING_EXTRACTIVE.items()}		,							info=_snake_case						)
									assert isinstance(dataset.info.task_templates		,							_snake_case						)
									assert len(dataset.info.task_templates						) == 1

									def keep_task(_snake_case						):
														return x

									def dont_keep_task(_snake_case						):
														_lowerCAmelCase         =       deepcopy(_snake_case						)
														_lowerCAmelCase         =       0
														return out

									_lowerCAmelCase         =       dataset.map(_snake_case						)
									assert mapped_dataset.info.task_templates == dataset.info.task_templates
									# reload from cache
									_lowerCAmelCase         =       dataset.map(_snake_case						)
									assert mapped_dataset.info.task_templates == dataset.info.task_templates

									_lowerCAmelCase         =       dataset.map(_snake_case						)
									assert mapped_dataset.info.task_templates == []
									# reload from cache
									_lowerCAmelCase         =       dataset.map(_snake_case						)
									assert mapped_dataset.info.task_templates == []




				def 				snake_case  (				self						):



									"""simple docstring"""





									_lowerCAmelCase         =       Features({"""text""": Value("""string"""						), """label""": ClassLabel(names=("""pos""", """neg""")						)}						)
									_lowerCAmelCase         =       TextClassification(text_column="""text"""		,							label_column="""label"""						)
									_lowerCAmelCase         =       DatasetInfo(features=_snake_case		,							task_templates=_snake_case						)
									_lowerCAmelCase         =       Dataset.from_dict({"""text""": ["""A sentence."""], """label""": ["""pos"""]}		,							info=_snake_case						)

									def process(_snake_case						):
														return example

									_lowerCAmelCase         =       dataset.remove_columns("""label"""						)
									_lowerCAmelCase         =       modified_dataset.map(_snake_case						)
									assert mapped_dataset.info.task_templates == []







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()


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



@unittest.skip('''The note-seq package currently throws an error on import'''	)
class      snake_case__	(   _A    ,							unittest.TestCase	):
				__lowerCAmelCase		=   SpectrogramDiffusionPipeline
				__lowerCAmelCase		=   PipelineTesterMixin.required_optional_params - {
				    '''callback''',
				    '''latents''',
				    '''callback_steps''',
				    '''output_type''',
				    '''num_images_per_prompt''',
				}
				__lowerCAmelCase		=   False
				__lowerCAmelCase		=   False
				__lowerCAmelCase		=   TOKENS_TO_AUDIO_GENERATION_PARAMS
				__lowerCAmelCase		=   TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										torch.manual_seed(0   )
										__magic_name__					:				Tuple				=   SpectrogramNotesEncoder(
										    max_length=2048  ,   vocab_size=1536  ,   d_model=768  ,   dropout_rate=0.1  ,   num_layers=1  ,   num_heads=1  ,   d_kv=4  ,   d_ff=2048  ,   feed_forward_proj='gated-gelu'  ,   )

										__magic_name__					:				Union[str, Any]				=   SpectrogramContEncoder(
										    input_dims=128  ,   targets_context_length=256  ,   d_model=768  ,   dropout_rate=0.1  ,   num_layers=1  ,   num_heads=1  ,   d_kv=4  ,   d_ff=2048  ,   feed_forward_proj='gated-gelu'  ,   )

										__magic_name__					:				Optional[Any]				=   TaFilmDecoder(
										    input_dims=128  ,   targets_length=256  ,   max_decoder_noise_time=20000.0  ,   d_model=768  ,   num_layers=1  ,   num_heads=1  ,   d_kv=4  ,   d_ff=2048  ,   dropout_rate=0.1  ,   )

										__magic_name__					:				Any				=   DDPMScheduler()

										__magic_name__					:				Tuple				=   {
										    'notes_encoder': notes_encoder.eval(),
										    'continuous_encoder': continuous_encoder.eval(),
										    'decoder': decoder.eval(),
										    'scheduler': scheduler,
										    'melgan': None,
										}
										return components
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase=0   ):




										"""simple docstring"""




										if str(_UpperCAmelCase   ).startswith('mps'   ):
																__magic_name__					:				str				=   torch.manual_seed(_UpperCAmelCase   )
										else:
																__magic_name__					:				Optional[int]				=   torch.Generator(device=_UpperCAmelCase   ).manual_seed(_UpperCAmelCase   )
										__magic_name__					:				List[str]				=   {
										    'input_tokens': [
										        [1134, 90, 1135, 1133, 1080, 112, 1132, 1080, 1133, 1079, 133, 1132, 1079, 1133, 1] + [0] * 2033
										    ],
										    'generator': generator,
										    'num_inference_steps': 4,
										    'output_type': 'mel',
										}
										return inputs
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Dict				=   'cpu'  # ensure determinism for the device-dependent torch.Generator
										__magic_name__					:				Optional[Any]				=   self.get_dummy_components()
										__magic_name__					:				Dict				=   SpectrogramDiffusionPipeline(**_UpperCAmelCase   )
										__magic_name__					:				Union[str, Any]				=   pipe.to(_UpperCAmelCase   )
										pipe.set_progress_bar_config(disable=_UpperCAmelCase   )

										__magic_name__					:				List[Any]				=   self.get_dummy_inputs(_UpperCAmelCase   )
										__magic_name__					:				Dict				=   pipe(**_UpperCAmelCase   )
										__magic_name__					:				Optional[Any]				=   output.audios

										__magic_name__					:				Dict				=   mel[0, -3:, -3:]

										assert mel_slice.shape == (3, 3)
										__magic_name__					:				Union[str, Any]				=   np.array(
										    [-11.512925, -4.788215, -0.46172905, -2.051715, -10.539147, -10.970963, -9.091634, 4.0, 4.0]   )
										assert np.abs(mel_slice.flatten() - expected_slice   ).max() < 1E-2
				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_save_load_local()
				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_dict_tuple_outputs_equivalent()
				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_save_load_optional_components()
				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_attention_slicing_forward_pass()
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										pass
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										pass





				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_progress_bar()






@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class      snake_case__	(   unittest.TestCase	):
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Tuple				=   torch_device

										__magic_name__					:				Optional[int]				=   SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
										__magic_name__					:				Dict				=   pipe.melgan
										__magic_name__					:				int				=   None

										__magic_name__					:				List[str]				=   pipe.to(_UpperCAmelCase   )
										pipe.set_progress_bar_config(disable=_UpperCAmelCase   )

										def callback(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																# decode mel to audio
																__magic_name__					:				Any				=   melgan(input_features=mel_output.astype(np.floataa   )   )[0]
																assert len(audio[0]   ) == 8_1920 * (step + 1)
																# simulate that audio is played
																return audio

										__magic_name__					:				Optional[Any]				=   MidiProcessor()
										__magic_name__					:				str				=   processor(_UpperCAmelCase   )

										__magic_name__					:				Optional[int]				=   input_tokens[:3]
										__magic_name__					:				Optional[Any]				=   torch.manual_seed(0   )
										pipe(_UpperCAmelCase  ,   num_inference_steps=5  ,   generator=_UpperCAmelCase  ,   callback=_UpperCAmelCase  ,   output_type='mel'   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Tuple				=   torch_device

										__magic_name__					:				str				=   SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
										__magic_name__					:				List[str]				=   pipe.to(_UpperCAmelCase   )
										pipe.set_progress_bar_config(disable=_UpperCAmelCase   )
										__magic_name__					:				Union[str, Any]				=   MidiProcessor()

										__magic_name__					:				List[Any]				=   processor(_UpperCAmelCase   )
										# just run two denoising loops
										__magic_name__					:				List[Any]				=   input_tokens[:2]

										__magic_name__					:				int				=   torch.manual_seed(0   )
										__magic_name__					:				int				=   pipe(_UpperCAmelCase  ,   num_inference_steps=2  ,   generator=_UpperCAmelCase   )

										__magic_name__					:				List[str]				=   output.audios[0]

										assert abs(np.abs(_UpperCAmelCase   ).sum() - 3612.841   ) < 1E-1





				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Any				=   torch_device

										__magic_name__					:				List[Any]				=   SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
										__magic_name__					:				Any				=   pipe.to(_UpperCAmelCase   )
										pipe.set_progress_bar_config(disable=_UpperCAmelCase   )

										__magic_name__					:				List[str]				=   MidiProcessor()

										__magic_name__					:				Union[str, Any]				=   processor(_UpperCAmelCase   )

										# just run 4 denoising loops
										__magic_name__					:				List[Any]				=   input_tokens[:4]

										__magic_name__					:				Dict				=   torch.manual_seed(0   )
										__magic_name__					:				Union[str, Any]				=   pipe(_UpperCAmelCase  ,   num_inference_steps=100  ,   generator=_UpperCAmelCase   )

										__magic_name__					:				Dict				=   output.audios[0]
										assert abs(np.abs(_UpperCAmelCase   ).sum() - 9389.1111   ) < 5E-2


