





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


_UpperCamelCase							=      logging.get_logger(__name__)
class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""
				__snake_case   :       Any	 =		["""input_features""", """is_longer"""]




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




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




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

											__lowerCamelCase      :      int         =mel[idx_front : idx_front + chunk_frames, :]
											__lowerCamelCase      :      Tuple         =mel[idx_middle : idx_middle + chunk_frames, :]
											__lowerCamelCase      :      List[Any]         =mel[idx_back : idx_back + chunk_frames, :]

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




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

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

																		if truncation == "fusion":
																									__lowerCamelCase      :      Optional[int]         =self._np_extract_fbank_features(__lowercase						,    self.mel_filters			)
																									__lowerCamelCase      :      Tuple         =np.stack([input_mel, input_mel, input_mel, input_mel]						,    axis=0			)
																		else:
																									__lowerCamelCase      :      Optional[int]         =self._np_extract_fbank_features(__lowercase						,    self.mel_filters_slaney			)[None, :]

											return input_mel, longer




				def __call__(     self  :List[str]						,    __lowercase  :Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]						,    __lowercase  :str = None						,    __lowercase  :Optional[str] = None						,    __lowercase  :Optional[int] = None						,    __lowercase  :Optional[int] = None						,    __lowercase  :Optional[Union[str, TensorType]] = None						,    **__lowercase  :Optional[Any]						,    ):
											__lowerCamelCase      :      List[str]         =truncation if truncation is not None else self.truncation
											__lowerCamelCase      :      Optional[int]         =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.'''			)

											__lowerCamelCase      :      List[Any]         =isinstance(__lowercase						,    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}'			)
											__lowerCamelCase      :      Union[str, Any]         =is_batched_numpy or (
											    isinstance(__lowercase						,    (list, tuple)			) and (isinstance(raw_speech[0]						,    (np.ndarray, tuple, list)			))
											)

											if is_batched:
																		__lowerCamelCase      :      Optional[Any]         =[np.asarray(__lowercase						,    dtype=np.floataa			) for speech in raw_speech]
											elif not is_batched and not isinstance(__lowercase						,    np.ndarray			):
																		__lowerCamelCase      :      Optional[Any]         =np.asarray(__lowercase						,    dtype=np.floataa			)
											elif isinstance(__lowercase						,    np.ndarray			) and raw_speech.dtype is np.dtype(np.floataa			):
																		__lowerCamelCase      :      int         =raw_speech.astype(np.floataa			)

											# always return batch
											if not is_batched:
																		__lowerCamelCase      :      Tuple         =[np.asarray(__lowercase			)]

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

											__lowerCamelCase      :      Any         =[]
											__lowerCamelCase      :      Dict         =[]
											for mel, longer in padded_inputs:
																		input_mel.append(__lowercase			)
																		is_longer.append(__lowercase			)

											if truncation == "fusion" and sum(__lowercase			) == 0:
																		# if no audio is longer than 10s, then randomly select one audio to be longer
																		__lowerCamelCase      :      Union[str, Any]         =np.random.randint(0						,    len(__lowercase			)			)
																		__lowerCamelCase      :      str         =True

											if isinstance(input_mel[0]						,    __lowercase			):
																		__lowerCamelCase      :      List[str]         =[np.asarray(__lowercase						,    dtype=np.floataa			) for feature in input_mel]

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

											__lowerCamelCase      :      Optional[Any]         ={'''input_features''': input_mel, '''is_longer''': is_longer}
											__lowerCamelCase      :      List[str]         =BatchFeature(__lowercase			)

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

											return input_features











"""simple docstring"""





from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


_UpperCamelCase							=      logging.get_logger(__name__)



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							str    ,						SCREAMING_SNAKE_CASE       :							str   ):


							'''simple docstring'''







							try:
														with open(SCREAMING_SNAKE_CASE    ,						'''rb'''   ) as flax_state_f:
																					__lowerCamelCase      :      int         =from_bytes(SCREAMING_SNAKE_CASE    ,						flax_state_f.read()   )
							except UnpicklingError as e:
														try:
																					with open(SCREAMING_SNAKE_CASE   ) as f:
																												if f.read().startswith('''version'''   ):
																																			raise OSError(
																																			    '''You seem to have cloned a repository without having git-lfs installed. Please'''
																																			    ''' install git-lfs and run `git lfs install` followed by `git lfs pull` in the'''
																																			    ''' folder you cloned.'''   )
																												else:
																																			raise ValueError from e
														except (UnicodeDecodeError, ValueError):
																					raise EnvironmentError(F'Unable to convert {model_file} to Flax deserializable object. '   )

							return load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							Union[str, Any]   ):


							'''simple docstring'''







							try:
														import torch  # noqa: F401
							except ImportError:
														logger.error(
														    '''Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see'''
														    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
														    ''' instructions.'''   )
														raise

							# check if we have bf16 weights
							__lowerCamelCase      :      str         =flatten_dict(jax.tree_util.tree_map(lambda SCREAMING_SNAKE_CASE   : x.dtype == jnp.bfloataa    ,						SCREAMING_SNAKE_CASE   )   ).values()
							if any(SCREAMING_SNAKE_CASE   ):
														# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

														# and bf16 is not fully supported in PT yet.
														logger.warning(
														    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
														    '''before loading those in PyTorch model.'''   )
														__lowerCamelCase      :      List[str]         =jax.tree_util.tree_map(
														    lambda SCREAMING_SNAKE_CASE   : params.astype(np.floataa   ) if params.dtype == jnp.bfloataa else params    ,						SCREAMING_SNAKE_CASE   )

							__lowerCamelCase      :      int         =''''''

							__lowerCamelCase      :      List[str]         =flatten_dict(SCREAMING_SNAKE_CASE    ,						sep='''.'''   )
							__lowerCamelCase      :      List[str]         =pt_model.state_dict()

							# keep track of unexpected & missing keys
							__lowerCamelCase      :      Union[str, Any]         =[]
							__lowerCamelCase      :      int         =set(pt_model_dict.keys()   )

							for flax_key_tuple, flax_tensor in flax_state_dict.items():
														__lowerCamelCase      :      Tuple         =flax_key_tuple.split('''.'''   )

														if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
																					__lowerCamelCase      :      Optional[int]         =flax_key_tuple_array[:-1] + ['''weight''']
																					__lowerCamelCase      :      Union[str, Any]         =jnp.transpose(SCREAMING_SNAKE_CASE    ,						(3, 2, 0, 1)   )
														elif flax_key_tuple_array[-1] == "kernel":
																					__lowerCamelCase      :      Optional[Any]         =flax_key_tuple_array[:-1] + ['''weight''']
																					__lowerCamelCase      :      Dict         =flax_tensor.T
														elif flax_key_tuple_array[-1] == "scale":
																					__lowerCamelCase      :      str         =flax_key_tuple_array[:-1] + ['''weight''']

														if "time_embedding" not in flax_key_tuple_array:
																					for i, flax_key_tuple_string in enumerate(SCREAMING_SNAKE_CASE   ):
																												__lowerCamelCase      :      Union[str, Any]         =(
																												    flax_key_tuple_string.replace('''_0'''    ,						'''.0'''   )
																												    .replace('''_1'''    ,						'''.1'''   )
																												    .replace('''_2'''    ,						'''.2'''   )
																												    .replace('''_3'''    ,						'''.3'''   )
																												    .replace('''_4'''    ,						'''.4'''   )
																												    .replace('''_5'''    ,						'''.5'''   )
																												    .replace('''_6'''    ,						'''.6'''   )
																												    .replace('''_7'''    ,						'''.7'''   )
																												    .replace('''_8'''    ,						'''.8'''   )
																												    .replace('''_9'''    ,						'''.9'''   )
																												)

														__lowerCamelCase      :      Optional[int]         ='''.'''.join(SCREAMING_SNAKE_CASE   )

														if flax_key in pt_model_dict:
																					if flax_tensor.shape != pt_model_dict[flax_key].shape:
																												raise ValueError(
																												    F'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
																												    F'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'   )
																					else:
																												# add weight to pytorch dict
																												__lowerCamelCase      :      Tuple         =np.asarray(SCREAMING_SNAKE_CASE   ) if not isinstance(SCREAMING_SNAKE_CASE    ,						np.ndarray   ) else flax_tensor
																												__lowerCamelCase      :      Any         =torch.from_numpy(SCREAMING_SNAKE_CASE   )
																												# remove from missing keys
																												missing_keys.remove(SCREAMING_SNAKE_CASE   )
														else:
																					# weight is not expected by PyTorch model
																					unexpected_keys.append(SCREAMING_SNAKE_CASE   )

							pt_model.load_state_dict(SCREAMING_SNAKE_CASE   )

							# re-transform missing_keys to list
							__lowerCamelCase      :      Optional[Any]         =list(SCREAMING_SNAKE_CASE   )

							if len(SCREAMING_SNAKE_CASE   ) > 0:
														logger.warning(
														    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
														    F' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
														    F' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
														    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
														    F' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
														    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
														    ''' FlaxBertForSequenceClassification model).'''   )
							if len(SCREAMING_SNAKE_CASE   ) > 0:
														logger.warning(
														    F'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
														    F' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
														    ''' use it for predictions and inference.'''   )

							return pt_model




