


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	: str  				=logging.get_logger(__name__)
class        _lowercase      (			A__       ):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			Tuple						=		['''input_features''', '''is_longer''']


				def __init__(       self			:Union[str, Any]			,			lowerCAmelCase__			:int=64			,			lowerCAmelCase__			:Tuple=48_000			,			lowerCAmelCase__			:List[str]=480			,			lowerCAmelCase__			:Optional[Any]=10			,			lowerCAmelCase__			:List[Any]=1_024			,			lowerCAmelCase__			:str=0.0			,			lowerCAmelCase__			:Tuple=False			,			lowerCAmelCase__			:float = 0			,			lowerCAmelCase__			:float = 14_000			,			lowerCAmelCase__			:int = None			,			lowerCAmelCase__			:str = "fusion"			,			lowerCAmelCase__			:str = "repeatpad"			,			**lowerCAmelCase__			:List[Any]			,			)		->			Tuple:
										super().__init__(
										    feature_size=lowerCAmelCase__			,			sampling_rate=lowerCAmelCase__			,			padding_value=lowerCAmelCase__			,			return_attention_mask=lowerCAmelCase__			,			**lowerCAmelCase__			,			)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		top_db
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		truncation
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		padding
										__SCREAMING_SNAKE_CASE						:     Any					=		fft_window_size
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		(fft_window_size >> 1) + 1
										__SCREAMING_SNAKE_CASE						:     Dict					=		hop_length
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		max_length_s
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		max_length_s * sampling_rate
										__SCREAMING_SNAKE_CASE						:     Any					=		sampling_rate
										__SCREAMING_SNAKE_CASE						:     Tuple					=		frequency_min
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		frequency_max
										__SCREAMING_SNAKE_CASE						:     int					=		mel_filter_bank(
										    num_frequency_bins=self.nb_frequency_bins			,			num_mel_filters=lowerCAmelCase__			,			min_frequency=lowerCAmelCase__			,			max_frequency=lowerCAmelCase__			,			sampling_rate=lowerCAmelCase__			,			norm=lowerCAmelCase__			,			mel_scale='''htk'''			,			)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		mel_filter_bank(
										    num_frequency_bins=self.nb_frequency_bins			,			num_mel_filters=lowerCAmelCase__			,			min_frequency=lowerCAmelCase__			,			max_frequency=lowerCAmelCase__			,			sampling_rate=lowerCAmelCase__			,			norm='''slaney'''			,			mel_scale='''slaney'''			,			)


				def 		__magic_name__(       self			:Any	)		->			Dict[str, Any]:
										__SCREAMING_SNAKE_CASE						:     str					=		copy.deepcopy(self.__dict__	)
										__SCREAMING_SNAKE_CASE						:     int					=		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 		__magic_name__(       self			:Optional[Any]			,			lowerCAmelCase__			:np.array			,			lowerCAmelCase__			:Optional[np.array] = None	)		->			np.ndarray:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		spectrogram(
										    lowerCAmelCase__			,			window_function(self.fft_window_size			,			'''hann'''	)			,			frame_length=self.fft_window_size			,			hop_length=self.hop_length			,			power=2.0			,			mel_filters=lowerCAmelCase__			,			log_mel='''dB'''			,			)
										return log_mel_spectrogram.T


				def 		__magic_name__(       self			:Dict			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:List[str]			,			lowerCAmelCase__			:List[str]	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		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						:     Any					=		[0]
										if len(ranges[2]	) == 0:
																# if the audio is too short, we just use the first chunk
																__SCREAMING_SNAKE_CASE						:     Tuple					=		[0]
										# randomly choose index for each part
										__SCREAMING_SNAKE_CASE						:     int					=		np.random.choice(ranges[0]	)
										__SCREAMING_SNAKE_CASE						:     Dict					=		np.random.choice(ranges[1]	)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		np.random.choice(ranges[2]	)

										__SCREAMING_SNAKE_CASE						:     List[str]					=		mel[idx_front : idx_front + chunk_frames, :]
										__SCREAMING_SNAKE_CASE						:     Tuple					=		mel[idx_middle : idx_middle + chunk_frames, :]
										__SCREAMING_SNAKE_CASE						:     List[str]					=		mel[idx_back : idx_back + chunk_frames, :]

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		torch.tensor(mel[None, None, :]	)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		torch.nn.functional.interpolate(
										    lowerCAmelCase__			,			size=[chunk_frames, 64]			,			mode='''bilinear'''			,			align_corners=lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		mel_shrink[0][0].numpy()
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]			,			axis=0	)
										return mel_fusion


				def 		__magic_name__(       self			:Union[str, Any]			,			lowerCAmelCase__			:np.array			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:Tuple	)		->			np.array:
										if waveform.shape[0] > max_length:
																if truncation == "rand_trunc":
																						__SCREAMING_SNAKE_CASE						:     int					=		True
																						# random crop to max_length (for compatibility) -> this should be handled by self.pad
																						__SCREAMING_SNAKE_CASE						:     List[Any]					=		len(lowerCAmelCase__	) - max_length
																						__SCREAMING_SNAKE_CASE						:     int					=		np.random.randint(0			,			overflow + 1	)
																						__SCREAMING_SNAKE_CASE						:     int					=		waveform[idx : idx + max_length]
																						__SCREAMING_SNAKE_CASE						:     List[Any]					=		self._np_extract_fbank_features(lowerCAmelCase__			,			self.mel_filters_slaney	)[None, :]
																elif truncation == "fusion":
																						__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		self._np_extract_fbank_features(lowerCAmelCase__			,			self.mel_filters	)
																						__SCREAMING_SNAKE_CASE						:     Tuple					=		max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
																						__SCREAMING_SNAKE_CASE						:     str					=		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						:     Dict					=		np.stack([mel, mel, mel, mel]			,			axis=0	)
																												__SCREAMING_SNAKE_CASE						:     int					=		False
																						else:
																												__SCREAMING_SNAKE_CASE						:     List[Any]					=		self._random_mel_fusion(lowerCAmelCase__			,			lowerCAmelCase__			,			lowerCAmelCase__	)
																												__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		True
																else:
																						raise NotImplementedError(f'''data_truncating {truncation} not implemented'''	)

										else:
																__SCREAMING_SNAKE_CASE						:     List[Any]					=		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						:     Optional[int]					=		int(max_length / len(lowerCAmelCase__	)	)
																												__SCREAMING_SNAKE_CASE						:     Dict					=		np.stack(np.tile(lowerCAmelCase__			,			n_repeat + 1	)	)[:max_length]
																						if padding == "repeatpad":
																												__SCREAMING_SNAKE_CASE						:     List[Any]					=		int(max_length / len(lowerCAmelCase__	)	)
																												__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		np.stack(np.tile(lowerCAmelCase__			,			lowerCAmelCase__	)	)
																						__SCREAMING_SNAKE_CASE						:     str					=		np.pad(lowerCAmelCase__			,			(0, max_length - waveform.shape[0])			,			mode='''constant'''			,			constant_values=0	)

																if truncation == "fusion":
																						__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self._np_extract_fbank_features(lowerCAmelCase__			,			self.mel_filters	)
																						__SCREAMING_SNAKE_CASE						:     Tuple					=		np.stack([input_mel, input_mel, input_mel, input_mel]			,			axis=0	)
																else:
																						__SCREAMING_SNAKE_CASE						:     Dict					=		self._np_extract_fbank_features(lowerCAmelCase__			,			self.mel_filters_slaney	)[None, :]

										return input_mel, longer



				def __call__(       self			:Tuple			,			lowerCAmelCase__			:Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]			,			lowerCAmelCase__			:str = None			,			lowerCAmelCase__			:Optional[str] = None			,			lowerCAmelCase__			:Optional[int] = None			,			lowerCAmelCase__			:Optional[int] = None			,			lowerCAmelCase__			:Optional[Union[str, TensorType]] = None			,			**lowerCAmelCase__			:Dict			,			)		->			BatchFeature:
										__SCREAMING_SNAKE_CASE						:     Tuple					=		truncation if truncation is not None else self.truncation
										__SCREAMING_SNAKE_CASE						:     Dict					=		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						:     Optional[Any]					=		isinstance(lowerCAmelCase__			,			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						:     Dict					=		is_batched_numpy or (
										    isinstance(lowerCAmelCase__			,			(list, tuple)	) and (isinstance(raw_speech[0]			,			(np.ndarray, tuple, list)	))
										)

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

										# always return batch
										if not is_batched:
																__SCREAMING_SNAKE_CASE						:     int					=		[np.asarray(lowerCAmelCase__	)]

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

										__SCREAMING_SNAKE_CASE						:     Dict					=		[]
										__SCREAMING_SNAKE_CASE						:     Tuple					=		[]
										for mel, longer in padded_inputs:
																input_mel.append(lowerCAmelCase__	)
																is_longer.append(lowerCAmelCase__	)

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

										if isinstance(input_mel[0]			,			lowerCAmelCase__	):
																__SCREAMING_SNAKE_CASE						:     int					=		[np.asarray(lowerCAmelCase__			,			dtype=np.floataa	) for feature in input_mel]

										# is_longer is a list of bool
										__SCREAMING_SNAKE_CASE						:     str					=		[[longer] for longer in is_longer]

										__SCREAMING_SNAKE_CASE						:     Dict					=		{'''input_features''': input_mel, '''is_longer''': is_longer}
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		BatchFeature(lowerCAmelCase__	)

										if return_tensors is not None:
																__SCREAMING_SNAKE_CASE						:     Optional[int]					=		input_features.convert_to_tensors(lowerCAmelCase__	)

										return input_features








import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm
							import torch_xla.debug.metrics as met
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				def __init__(     self						:    Dict,  *lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    Any=None,  **lowerCamelCase						:    str  ):

								'''simple docstring'''
								super().__init__(*lowerCamelCase,  **lowerCamelCase  )
								lowercase__    =			eval_examples
								lowercase__    =			post_process_function




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    str=None,  lowerCamelCase						:    Optional[Any]=None,  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    str = "eval"  ):

								'''simple docstring'''
								lowercase__    =			self.eval_dataset if eval_dataset is None else eval_dataset
								lowercase__    =			self.get_eval_dataloader(lowerCamelCase  )
								lowercase__    =			self.eval_examples if eval_examples is None else eval_examples

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Evaluation''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )
								if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
												# Only the main node write the results by default
												lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions  )
												lowercase__    =			self.compute_metrics(lowerCamelCase  )

												# Prefix all keys with metric_key_prefix + '_'
												for key in list(metrics.keys()  ):
																if not key.startswith(F"""{metric_key_prefix}_"""  ):
																				lowercase__    =			metrics.pop(lowerCamelCase  )
												metrics.update(output.metrics  )
								else:
												lowercase__    =			output.metrics

								if self.args.should_log:
												# Only the main node log the results by default
												self.log(lowerCamelCase  )

								if self.args.tpu_metrics_debug or self.args.debug:
												# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
												xm.master_print(met.metrics_report()  )

								lowercase__    =			self.callback_handler.on_evaluate(self.args,  self.state,  self.control,  lowerCamelCase  )
								return metrics




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    Any,  lowerCamelCase						:    Dict,  lowerCamelCase						:    int=None,  lowerCamelCase						:    str = "test"  ):

								'''simple docstring'''
								lowercase__    =			self.get_test_dataloader(lowerCamelCase  )

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Prediction''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )

								if self.post_process_function is None or self.compute_metrics is None:
												return output

								lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions,  '''predict'''  )
								lowercase__    =			self.compute_metrics(lowerCamelCase  )

								# Prefix all keys with metric_key_prefix + '_'
								for key in list(metrics.keys()  ):
												if not key.startswith(F"""{metric_key_prefix}_"""  ):
																lowercase__    =			metrics.pop(lowerCamelCase  )
								metrics.update(output.metrics  )
								return PredictionOutput(predictions=predictions.predictions,  label_ids=predictions.label_ids,  metrics=lowerCamelCase  )




