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


lowerCamelCase_				:		List[Any]       						=       """./tests/fixtures/elise_format0.mid"""


@unittest.skip('The note-seq package currently throws an error on import'		)
class 						a__							(   __snake_case      ,						unittest.TestCase		):
			A__				:   Tuple								=			SpectrogramDiffusionPipeline
			A__				:   Any								=			PipelineTesterMixin.required_optional_params - {
			    'callback',
			    'latents',
			    'callback_steps',
			    'output_type',
			    'num_images_per_prompt',
			}
			A__				:   List[Any]								=			False
			A__				:   int								=			False
			A__				:   Any								=			TOKENS_TO_AUDIO_GENERATION_PARAMS
			A__				:   Optional[int]								=			TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Dict:
								torch.manual_seed(0    )
								__a							=   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'      ,						)

								__a							=   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'      ,						)

								__a							=   TaFilmDecoder(
								    input_dims=1_2_8      ,						targets_length=2_5_6      ,						max_decoder_noise_time=20_000.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      ,						)

								__a							=   DDPMScheduler()

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



			def __SCREAMING_SNAKE_CASE							(						self      ,						UpperCAmelCase      ,						UpperCAmelCase=0    ) ->    int:
								if str(UpperCAmelCase    ).startswith('mps'    ):
													__a							=   torch.manual_seed(UpperCAmelCase    )
								else:
													__a							=   torch.Generator(device=UpperCAmelCase    ).manual_seed(UpperCAmelCase    )
								__a							=   {
								    '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 __SCREAMING_SNAKE_CASE							(						self    ) ->    str:
								__a							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
								__a							=   self.get_dummy_components()
								__a							=   SpectrogramDiffusionPipeline(**UpperCAmelCase    )
								__a							=   pipe.to(UpperCAmelCase    )
								pipe.set_progress_bar_config(disable=UpperCAmelCase    )

								__a							=   self.get_dummy_inputs(UpperCAmelCase    )
								__a							=   pipe(**UpperCAmelCase    )
								__a							=   output.audios

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

								assert mel_slice.shape == (3, 3)
								__a							=   np.array(
								    [-11.512_925, -4.788_215, -0.46_172_905, -2.051_715, -10.539_147, -10.970_963, -9.091_634, 4.0, 4.0]    )
								assert np.abs(mel_slice.flatten() - expected_slice    ).max() < 1e-2



			@skip_mps
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Union[str, Any]:
								return super().test_save_load_local()



			@skip_mps
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Optional[int]:
								return super().test_dict_tuple_outputs_equivalent()



			@skip_mps
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Optional[Any]:
								return super().test_save_load_optional_components()



			@skip_mps
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    str:
								return super().test_attention_slicing_forward_pass()



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    int:
								pass



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    str:
								pass







			@skip_mps
			def __SCREAMING_SNAKE_CASE							(						self    ) ->    int:
								return super().test_progress_bar()






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



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Any:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Tuple:
								# TODO - test that pipeline can decode tokens in a callback
								# so that music can be played live
								__a							=   torch_device

								__a							=   SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'    )
								__a							=   pipe.melgan
								__a							=   None

								__a							=   pipe.to(UpperCAmelCase    )
								pipe.set_progress_bar_config(disable=UpperCAmelCase    )

								def callback(UpperCAmelCase      ,						UpperCAmelCase    ):
													# decode mel to audio
													__a							=   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

								__a							=   MidiProcessor()
								__a							=   processor(UpperCAmelCase    )

								__a							=   input_tokens[:3]
								__a							=   torch.manual_seed(0    )
								pipe(UpperCAmelCase      ,						num_inference_steps=5      ,						generator=UpperCAmelCase      ,						callback=UpperCAmelCase      ,						output_type='mel'    )



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Tuple:
								__a							=   torch_device

								__a							=   SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'    )
								__a							=   pipe.to(UpperCAmelCase    )
								pipe.set_progress_bar_config(disable=UpperCAmelCase    )
								__a							=   MidiProcessor()

								__a							=   processor(UpperCAmelCase    )
								# just run two denoising loops
								__a							=   input_tokens[:2]

								__a							=   torch.manual_seed(0    )
								__a							=   pipe(UpperCAmelCase      ,						num_inference_steps=2      ,						generator=UpperCAmelCase    )

								__a							=   output.audios[0]

								assert abs(np.abs(UpperCAmelCase    ).sum() - 3_612.841    ) < 1e-1







			def __SCREAMING_SNAKE_CASE							(						self    ) ->    List[str]:
								__a							=   torch_device

								__a							=   SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'    )
								__a							=   pipe.to(UpperCAmelCase    )
								pipe.set_progress_bar_config(disable=UpperCAmelCase    )

								__a							=   MidiProcessor()

								__a							=   processor(UpperCAmelCase    )

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

								__a							=   torch.manual_seed(0    )
								__a							=   pipe(UpperCAmelCase      ,						num_inference_steps=1_0_0      ,						generator=UpperCAmelCase    )

								__a							=   output.audios[0]
								assert abs(np.abs(UpperCAmelCase    ).sum() - 9_389.1_111    ) < 5e-2





from typing import List

from packaging import version
from sklearn.metrics import fa_score

import datasets
from datasets.config import PY_VERSION


if PY_VERSION < version.parse("""3.8"""):
	import importlib_metadata
else:
	import importlib.metadata as importlib_metadata


# TODO(Patrick/Anton)
lowerCamelCase_				:		List[str]       						=       """\
"""

lowerCamelCase_				:		Any       						=       """\
    XTREME-S is a benchmark to evaluate universal cross-lingual speech representations in many languages.
    XTREME-S covers four task families: speech recognition, classification, speech-to-text translation and retrieval.
"""

lowerCamelCase_				:		Tuple       						=       """
Compute XTREME-S evaluation metric associated to each XTREME-S dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
    bleu_kwargs: optional dict of keywords to be passed when computing 'bleu'.
        Keywords include Dict can be one of 'smooth_method', 'smooth_value', 'force', 'lowercase',
        'tokenize', 'use_effective_order'.
    wer_kwargs: optional dict of keywords to be passed when computing 'wer' and 'cer'.
        Keywords include 'concatenate_texts'.
Returns: depending on the XTREME-S task, one or several of:
    \"accuracy\": Accuracy - for 'fleurs-lang_id', 'minds14'
    \"f1\": F1 score - for 'minds14'
    \"wer\": Word error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'
    \"cer\": Character error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'
    \"bleu\": BLEU score according to the `sacrebleu` metric - for 'covost2'
Examples:

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'mls')  # 'mls', 'voxpopuli', 'fleurs-asr' or 'babel'
    >>> references = [\"it is sunny here\", \"paper and pen are essentials\"]
    >>> predictions = [\"it's sunny\", \"paper pen are essential\"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'wer': 0.56, 'cer': 0.27}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'covost2')
    >>> references = [\"bonjour paris\", \"il est necessaire de faire du sport de temps en temp\"]
    >>> predictions = [\"bonjour paris\", \"il est important de faire du sport souvent\"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'bleu': 31.65}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'fleurs-lang_id')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'accuracy': 0.6}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'minds14')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'f1': 0.58, 'accuracy': 0.6}
"""

lowerCamelCase_				:		Dict       						=       ["""fleurs-asr""", """mls""", """voxpopuli""", """babel""", """covost2""", """fleurs-lang_id""", """minds14"""]
lowerCamelCase_				:		Optional[Any]       						=       """"""

try:
	from jiwer import transforms as tr

	lowerCamelCase_				:		str       						=       True
except ImportError:
	lowerCamelCase_				:		List[Any]       						=       False

if _jiwer_available and version.parse(importlib_metadata.version("""jiwer""")) < version.parse("""2.3.0"""):


	class 						a__							(   tr.AbstractTransform		):



				def __init__(						self      ,						UpperCAmelCase = " "    ) ->    List[str]:
									__a							=   sentence_delimiter



				def __SCREAMING_SNAKE_CASE							(						self      ,						UpperCAmelCase    ) ->    List[str]:
									return list(UpperCAmelCase    )







				def __SCREAMING_SNAKE_CASE							(						self      ,						UpperCAmelCase    ) ->    Optional[int]:
									__a							=   []
									for sent_idx, sentence in enumerate(UpperCAmelCase    ):
														chars.extend(self.process_string(UpperCAmelCase    )    )
														if self.sentence_delimiter is not None and self.sentence_delimiter != "" and sent_idx < len(UpperCAmelCase    ) - 1:
																			chars.append(self.sentence_delimiter    )
									return chars






	lowerCamelCase_				:		Any       						=       tr.Compose(
	    [tr.RemoveMultipleSpaces(), tr.Strip(), SentencesToListOfCharacters(SENTENCE_DELIMITER)]
	)
elif _jiwer_available:
	lowerCamelCase_				:		Optional[Any]       						=       tr.Compose(
	    [
	        tr.RemoveMultipleSpaces(),
	        tr.Strip(),
	        tr.ReduceToSingleSentence(SENTENCE_DELIMITER),
	        tr.ReduceToListOfListOfChars(),
	    ]
	)
else:
	lowerCamelCase_				:		Tuple       						=       None
def 							lowerCAmelCase(	__lowerCamelCase      ,				__lowerCamelCase						):
					return float((preds == labels).mean()						)
def 							lowerCAmelCase(	__lowerCamelCase      ,				__lowerCamelCase						):
					return {
					    "f1": float(fa_score(y_true=__lowerCamelCase      ,				y_pred=__lowerCamelCase      ,				average='macro'						)						),
					    "accuracy": simple_accuracy(__lowerCamelCase      ,				__lowerCamelCase						),
					}
def 							lowerCAmelCase(	__lowerCamelCase      ,				__lowerCamelCase      ,				__lowerCamelCase="exp"      ,				__lowerCamelCase=None      ,				__lowerCamelCase=False      ,				__lowerCamelCase=False      ,				__lowerCamelCase=None      ,				__lowerCamelCase=False      ,				):
					# xtreme-s can only have one label
					__a							=   [[label] for label in labels]
					__a							=   list(__lowerCamelCase						)
					try:
										import sacrebleu as scb
					except ImportError:
										raise ValueError(
										    'sacrebleu has to be installed in order to apply the bleu metric for covost2.'
										    'You can install it via `pip install sacrebleu`.'						)

					if version.parse(scb.__version__						) < version.parse('1.4.12'						):
										raise ImportWarning(
										    'To use `sacrebleu`, the module `sacrebleu>=1.4.12` is required, and the current version of `sacrebleu` doesn\'t match this condition.\n'
										    'You can install it with `pip install "sacrebleu>=1.4.12"`.'						)

					__a							=   len(labels[0]						)
					if any(len(__lowerCamelCase						) != references_per_prediction for refs in labels						):
										raise ValueError('Sacrebleu requires the same number of references for each prediction'						)
					__a							=   [[refs[i] for refs in labels] for i in range(__lowerCamelCase						)]
					__a							=   scb.corpus_bleu(
					    __lowerCamelCase      ,				__lowerCamelCase      ,				smooth_method=__lowerCamelCase      ,				smooth_value=__lowerCamelCase      ,				force=__lowerCamelCase      ,				lowercase=__lowerCamelCase      ,				use_effective_order=__lowerCamelCase      ,				**({'tokenize': tokenize} if tokenize else {})      ,				)
					return {"bleu": output.score}
def 							lowerCAmelCase(	__lowerCamelCase      ,				__lowerCamelCase      ,				__lowerCamelCase      ,				__lowerCamelCase						):
					try:
										from jiwer import compute_measures
					except ImportError:
										raise ValueError(
										    f'''jiwer has to be installed in order to apply the wer metric for {config_name}.'''
										    'You can install it via `pip install jiwer`.'						)

					if concatenate_texts:
										__a							=   compute_measures(__lowerCamelCase      ,				__lowerCamelCase						)['wer']

										__a							=   compute_measures(__lowerCamelCase      ,				__lowerCamelCase      ,				truth_transform=__lowerCamelCase      ,				hypothesis_transform=__lowerCamelCase						)['wer']
										return {"wer": wer, "cer": cer}
					else:

										def compute_score(__lowerCamelCase      ,				__lowerCamelCase      ,				__lowerCamelCase="wer"						):
															__a							=   0
															__a							=   0
															for prediction, reference in zip(__lowerCamelCase      ,				__lowerCamelCase						):
																				if score_type == "wer":
																									__a							=   compute_measures(__lowerCamelCase      ,				__lowerCamelCase						)
																				elif score_type == "cer":
																									__a							=   compute_measures(
																									    __lowerCamelCase      ,				__lowerCamelCase      ,				truth_transform=__lowerCamelCase      ,				hypothesis_transform=__lowerCamelCase						)
																				incorrect += measures["substitutions"] + measures["deletions"] + measures["insertions"]
																				total += measures["substitutions"] + measures["deletions"] + measures["hits"]
															return incorrect / total

										return {"wer": compute_score(__lowerCamelCase      ,				__lowerCamelCase      ,				'wer'						), "cer": compute_score(__lowerCamelCase      ,				__lowerCamelCase      ,				'cer'						)}



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION      ,						_KWARGS_DESCRIPTION		)
class 						a__							(   datasets.Metric		):



			def __SCREAMING_SNAKE_CASE							(						self    ) ->    Tuple:
								if self.config_name not in _CONFIG_NAMES:
													raise KeyError(f'''You should supply a configuration name selected in {_CONFIG_NAMES}'''    )

								__a							=   'int64' if self.config_name in ['fleurs-lang_id', 'minds14'] else 'string'

								return datasets.MetricInfo(
								    description=_DESCRIPTION      ,						citation=_CITATION      ,						inputs_description=_KWARGS_DESCRIPTION      ,						features=datasets.Features(
								        {'predictions': datasets.Value(UpperCAmelCase    ), 'references': datasets.Value(UpperCAmelCase    )}    )      ,						codebase_urls=[]      ,						reference_urls=[]      ,						format='numpy'      ,						)







			def __SCREAMING_SNAKE_CASE							(						self      ,						UpperCAmelCase      ,						UpperCAmelCase      ,						UpperCAmelCase=None      ,						UpperCAmelCase=None    ) ->    Optional[Any]:
								__a							=   bleu_kwargs if bleu_kwargs is not None else {}
								__a							=   wer_kwargs if wer_kwargs is not None else {}

								if self.config_name == "fleurs-lang_id":
													return {"accuracy": simple_accuracy(UpperCAmelCase      ,						UpperCAmelCase    )}
								elif self.config_name == "minds14":
													return fa_and_simple_accuracy(UpperCAmelCase      ,						UpperCAmelCase    )
								elif self.config_name == "covost2":
													__a							=   bleu_kwargs.pop('smooth_method'      ,						'exp'    )
													__a							=   bleu_kwargs.pop('smooth_value'      ,						UpperCAmelCase    )
													__a							=   bleu_kwargs.pop('force'      ,						UpperCAmelCase    )
													__a							=   bleu_kwargs.pop('lowercase'      ,						UpperCAmelCase    )
													__a							=   bleu_kwargs.pop('tokenize'      ,						UpperCAmelCase    )
													__a							=   bleu_kwargs.pop('use_effective_order'      ,						UpperCAmelCase    )
													return bleu(
													    preds=UpperCAmelCase      ,						labels=UpperCAmelCase      ,						smooth_method=UpperCAmelCase      ,						smooth_value=UpperCAmelCase      ,						force=UpperCAmelCase      ,						lowercase=UpperCAmelCase      ,						tokenize=UpperCAmelCase      ,						use_effective_order=UpperCAmelCase      ,						)
								elif self.config_name in ["fleurs-asr", "mls", "voxpopuli", "babel"]:
													__a							=   wer_kwargs.pop('concatenate_texts'      ,						UpperCAmelCase    )
													return wer_and_cer(UpperCAmelCase      ,						UpperCAmelCase      ,						UpperCAmelCase      ,						self.config_name    )
								else:
													raise KeyError(f'''You should supply a configuration name selected in {_CONFIG_NAMES}'''    )




