

import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np

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    :       Optional[int]        =							logging.get_logger(__name__)





class  __UpperCAmelCase					(	lowerCamelCase__						):
   UpperCamelCase       					=					["""input_features"""]



   def __init__(							self							:					Any,		__A							:					Optional[int]=8_0,		__A							:					List[str]=1_6_0_0_0,		__A							:					Any=1_6_0,		__A							:					Union[str, Any]=3_0,		__A							:					Optional[Any]=4_0_0,		__A							:					Optional[int]=0.0,		__A							:					Optional[Any]=False,		**__A							:					List[Any],		):
     super().__init__(
         feature_size=__A,		sampling_rate=__A,		padding_value=__A,		return_attention_mask=__A,		**__A,		)
     UpperCAmelCase	:							int						=			n_fft
     UpperCAmelCase	:							Optional[int]						=			hop_length
     UpperCAmelCase	:							Dict						=			chunk_length
     UpperCAmelCase	:							List[Any]						=			chunk_length * sampling_rate
     UpperCAmelCase	:							Dict						=			self.n_samples // hop_length
     UpperCAmelCase	:							Any						=			sampling_rate
     UpperCAmelCase	:							Dict						=			mel_filter_bank(
         num_frequency_bins=1 + n_fft // 2,		num_mel_filters=__A,		min_frequency=0.0,		max_frequency=8_0_0_0.0,		sampling_rate=__A,		norm='''slaney''',		mel_scale='''slaney''',		)



   def 	__magic_name__   (							self							:					Tuple,		__A							:					np.array						):
     UpperCAmelCase	:							List[str]						=			spectrogram(
         __A,		window_function(self.n_fft,		'''hann'''						),		frame_length=self.n_fft,		hop_length=self.hop_length,		power=2.0,		mel_filters=self.mel_filters,		log_mel='''log10''',		)
     UpperCAmelCase	:							List[Any]						=			log_spec[:, :-1]
     UpperCAmelCase	:							Optional[Any]						=			np.maximum(__A,		log_spec.max() - 8.0						)
     UpperCAmelCase	:							Any						=			(log_spec + 4.0) / 4.0
     return log_spec



   @staticmethod
   # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
   def 	__magic_name__   (							__A							:					List[np.ndarray],		__A							:					List[np.ndarray],		__A							:					float = 0.0						):
     if attention_mask is not None:
       UpperCAmelCase	:							Tuple						=			np.array(__A,		np.intaa						)
       UpperCAmelCase	:							str						=			[]

       for vector, length in zip(__A,		attention_mask.sum(-1						)						):
         UpperCAmelCase	:							Optional[Any]						=			(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1E-7						)
         if length < normed_slice.shape[0]:
           UpperCAmelCase	:							Any						=			padding_value

         normed_input_values.append(__A						)
     else:
       UpperCAmelCase	:							int						=			[(x - x.mean()) / np.sqrt(x.var() + 1E-7						) for x in input_values]

     return normed_input_values



   def __call__(							self							:					Any,		__A							:					Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],		__A							:					bool = True,		__A							:					Optional[int] = None,		__A							:					Optional[Union[str, TensorType]] = None,		__A							:					Optional[bool] = None,		__A							:					Optional[str] = "max_length",		__A							:					Optional[int] = None,		__A							:					Optional[int] = None,		__A							:					Optional[bool] = None,		**__A							:					Tuple,		):

     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.'''						)

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

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

     # always return batch
     if not is_batched:
       UpperCAmelCase	:							str						=			[np.asarray([raw_speech]						).T]

     UpperCAmelCase	:							Any						=			BatchFeature({'''input_features''': raw_speech}						)

     # convert into correct format for padding

     UpperCAmelCase	:							Optional[Any]						=			self.pad(
         __A,		padding=__A,		max_length=max_length if max_length else self.n_samples,		truncation=__A,		pad_to_multiple_of=__A,		return_attention_mask=return_attention_mask or do_normalize,		)

     # zero-mean and unit-variance normalization
     if do_normalize:
       UpperCAmelCase	:							Any						=			self.zero_mean_unit_var_norm(
           padded_inputs['''input_features'''],		attention_mask=padded_inputs['''attention_mask'''],		padding_value=self.padding_value,		)
       UpperCAmelCase	:							Union[str, Any]						=			np.stack(padded_inputs['''input_features'''],		axis=0						)

     # make sure list is in array format
     UpperCAmelCase	:							Union[str, Any]						=			padded_inputs.get('''input_features'''						).transpose(2,		0,		1						)

     UpperCAmelCase	:							Dict						=			[self._np_extract_fbank_features(__A						) for waveform in input_features[0]]

     if isinstance(input_features[0],		__A						):
       UpperCAmelCase	:							Optional[Any]						=			[np.asarray(__A,		dtype=np.floataa						) for feature in input_features]
     else:
       UpperCAmelCase	:							int						=			input_features

     if return_attention_mask:
       # rescale from sample (48000) to feature (3000)
       UpperCAmelCase	:							Optional[Any]						=			padded_inputs['''attention_mask'''][:, :: self.hop_length]

     if return_tensors is not None:
       UpperCAmelCase	:							List[str]						=			padded_inputs.convert_to_tensors(__A						)

     return padded_inputs





   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							List[str]						=			copy.deepcopy(self.__dict__						)
     UpperCAmelCase	:							Any						=			self.__class__.__name__
     if "mel_filters" in output:
       del output["mel_filters"]
     return output








import torch
from torch import nn





class  __UpperCAmelCase					(	nn.Module						):



   def __init__(							self							:					List[Any],		__A							:					List[Any],		__A							:					Optional[Any],		__A							:					int,		__A							:					List[Any],		__A							:					int=1,		__A							:					List[str]=False						):
     super().__init__()

     UpperCAmelCase	:							Union[str, Any]						=			n_token
     UpperCAmelCase	:							List[str]						=			d_embed
     UpperCAmelCase	:							Dict						=			d_proj

     UpperCAmelCase	:							List[Any]						=			cutoffs + [n_token]
     UpperCAmelCase	:							Dict						=			[0] + self.cutoffs
     UpperCAmelCase	:							int						=			div_val

     UpperCAmelCase	:							Union[str, Any]						=			self.cutoffs[0]
     UpperCAmelCase	:							str						=			len(self.cutoffs						) - 1
     UpperCAmelCase	:							Optional[int]						=			self.shortlist_size + self.n_clusters

     if self.n_clusters > 0:
       UpperCAmelCase	:							str						=			nn.Parameter(torch.zeros(self.n_clusters,		self.d_embed						)						)
       UpperCAmelCase	:							List[str]						=			nn.Parameter(torch.zeros(self.n_clusters						)						)

     UpperCAmelCase	:							Dict						=			nn.ModuleList()
     UpperCAmelCase	:							Optional[int]						=			nn.ParameterList()

     if div_val == 1:
       for i in range(len(self.cutoffs						)						):
         if d_proj != d_embed:
           self.out_projs.append(nn.Parameter(torch.FloatTensor(__A,		__A						)						)						)
         else:
           self.out_projs.append(__A						)

       self.out_layers.append(nn.Linear(__A,		__A						)						)
     else:
       for i in range(len(self.cutoffs						)						):
         UpperCAmelCase			,						UpperCAmelCase	:							Dict						=			self.cutoff_ends[i], self.cutoff_ends[i + 1]
         UpperCAmelCase	:							str						=			d_embed // (div_val**i)

         self.out_projs.append(nn.Parameter(torch.FloatTensor(__A,		__A						)						)						)

         self.out_layers.append(nn.Linear(__A,		r_idx - l_idx						)						)

     UpperCAmelCase	:							Optional[int]						=			keep_order



   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					List[str],		__A							:					Any,		__A							:					Dict,		__A							:					Optional[Any]						):
     if proj is None:
       UpperCAmelCase	:							List[Any]						=			nn.functional.linear(__A,		__A,		bias=__A						)
     else:
       # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
       UpperCAmelCase	:							Union[str, Any]						=			nn.functional.linear(__A,		proj.t().contiguous()						)
       UpperCAmelCase	:							Optional[int]						=			nn.functional.linear(__A,		__A,		bias=__A						)
       # else:
       #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
       #     if bias is not None:
       #         logit = logit + bias

     return logit



   def 	__magic_name__   (							self							:					int,		__A							:					int,		__A							:					List[Any]=None,		__A							:					Dict=False						):

     if labels is not None:
       # Shift so that tokens < n predict n
       UpperCAmelCase	:							List[Any]						=			hidden[..., :-1, :].contiguous()
       UpperCAmelCase	:							Any						=			labels[..., 1:].contiguous()
       UpperCAmelCase	:							Optional[Any]						=			hidden.view(-1,		hidden.size(-1						)						)
       UpperCAmelCase	:							Union[str, Any]						=			labels.view(-1						)
       if hidden.size(0						) != labels.size(0						):
         raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''						)
     else:
       UpperCAmelCase	:							str						=			hidden.view(-1,		hidden.size(-1						)						)

     if self.n_clusters == 0:
       UpperCAmelCase	:							List[str]						=			self._compute_logit(__A,		self.out_layers[0].weight,		self.out_layers[0].bias,		self.out_projs[0]						)
       if labels is not None:
         UpperCAmelCase	:							Optional[int]						=			labels != -1_0_0
         UpperCAmelCase	:							Dict						=			torch.zeros_like(__A,		dtype=hidden.dtype,		device=hidden.device						)
         UpperCAmelCase	:							Any						=			(
             -nn.functional.log_softmax(__A,		dim=-1						)[mask].gather(1,		labels[mask].unsqueeze(1						)						).squeeze(1						)
         )
       else:
         UpperCAmelCase	:							Any						=			nn.functional.log_softmax(__A,		dim=-1						)
     else:
       # construct weights and biases
       UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			[], []
       for i in range(len(self.cutoffs						)						):
         if self.div_val == 1:
           UpperCAmelCase			,						UpperCAmelCase	:							List[Any]						=			self.cutoff_ends[i], self.cutoff_ends[i + 1]
           UpperCAmelCase	:							List[str]						=			self.out_layers[0].weight[l_idx:r_idx]
           UpperCAmelCase	:							Optional[int]						=			self.out_layers[0].bias[l_idx:r_idx]
         else:
           UpperCAmelCase	:							List[Any]						=			self.out_layers[i].weight
           UpperCAmelCase	:							Dict						=			self.out_layers[i].bias

         if i == 0:
           UpperCAmelCase	:							List[str]						=			torch.cat([weight_i, self.cluster_weight],		dim=0						)
           UpperCAmelCase	:							List[Any]						=			torch.cat([bias_i, self.cluster_bias],		dim=0						)

         weights.append(__A						)
         biases.append(__A						)

       UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							str						=			weights[0], biases[0], self.out_projs[0]

       UpperCAmelCase	:							Dict						=			self._compute_logit(__A,		__A,		__A,		__A						)
       UpperCAmelCase	:							int						=			nn.functional.log_softmax(__A,		dim=1						)

       if labels is None:
         UpperCAmelCase	:							Optional[int]						=			hidden.new_empty((head_logit.size(0						), self.n_token)						)
       else:
         UpperCAmelCase	:							Union[str, Any]						=			torch.zeros_like(__A,		dtype=hidden.dtype,		device=hidden.device						)

       UpperCAmelCase	:							Union[str, Any]						=			0
       UpperCAmelCase	:							Any						=			[0] + self.cutoffs
       for i in range(len(__A						) - 1						):
         UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			cutoff_values[i], cutoff_values[i + 1]

         if labels is not None:
           UpperCAmelCase	:							List[str]						=			(labels >= l_idx) & (labels < r_idx)
           UpperCAmelCase	:							Tuple						=			mask_i.nonzero().squeeze()

           if indices_i.numel() == 0:
             continue

           UpperCAmelCase	:							Any						=			labels.index_select(0,		__A						) - l_idx
           UpperCAmelCase	:							Dict						=			head_logprob.index_select(0,		__A						)
           UpperCAmelCase	:							List[str]						=			hidden.index_select(0,		__A						)
         else:
           UpperCAmelCase	:							Tuple						=			hidden

         if i == 0:
           if labels is not None:
             UpperCAmelCase	:							Union[str, Any]						=			head_logprob_i.gather(1,		target_i[:, None]						).squeeze(1						)
           else:
             UpperCAmelCase	:							Optional[int]						=			head_logprob[:, : self.cutoffs[0]]
         else:
           UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							List[str]						=			weights[i], biases[i], self.out_projs[i]

           UpperCAmelCase	:							List[str]						=			self._compute_logit(__A,		__A,		__A,		__A						)
           UpperCAmelCase	:							Dict						=			nn.functional.log_softmax(__A,		dim=1						)
           UpperCAmelCase	:							int						=			self.cutoffs[0] + i - 1  # No probability for the head cluster
           if labels is not None:
             UpperCAmelCase	:							Union[str, Any]						=			head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                 1,		target_i[:, None]						).squeeze(1						)
           else:
             UpperCAmelCase	:							int						=			head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
             UpperCAmelCase	:							Optional[Any]						=			logprob_i

         if labels is not None:
           if (hasattr(self,		'''keep_order'''						) and self.keep_order) or keep_order:
             out.index_copy_(0,		__A,		-logprob_i						)
           else:
             out[offset : offset + logprob_i.size(0						)].copy_(-logprob_i						)
           offset += logprob_i.size(0						)

     return out





   def 	__magic_name__   (							self							:					Tuple,		__A							:					List[str]						):
     if self.n_clusters == 0:
       UpperCAmelCase	:							int						=			self._compute_logit(__A,		self.out_layers[0].weight,		self.out_layers[0].bias,		self.out_projs[0]						)
       return nn.functional.log_softmax(__A,		dim=-1						)
     else:
       # construct weights and biases
       UpperCAmelCase			,						UpperCAmelCase	:							Any						=			[], []
       for i in range(len(self.cutoffs						)						):
         if self.div_val == 1:
           UpperCAmelCase			,						UpperCAmelCase	:							Any						=			self.cutoff_ends[i], self.cutoff_ends[i + 1]
           UpperCAmelCase	:							List[Any]						=			self.out_layers[0].weight[l_idx:r_idx]
           UpperCAmelCase	:							List[str]						=			self.out_layers[0].bias[l_idx:r_idx]
         else:
           UpperCAmelCase	:							List[str]						=			self.out_layers[i].weight
           UpperCAmelCase	:							Dict						=			self.out_layers[i].bias

         if i == 0:
           UpperCAmelCase	:							Dict						=			torch.cat([weight_i, self.cluster_weight],		dim=0						)
           UpperCAmelCase	:							str						=			torch.cat([bias_i, self.cluster_bias],		dim=0						)

         weights.append(__A						)
         biases.append(__A						)

       UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Tuple						=			weights[0], biases[0], self.out_projs[0]
       UpperCAmelCase	:							int						=			self._compute_logit(__A,		__A,		__A,		__A						)

       UpperCAmelCase	:							Optional[int]						=			hidden.new_empty((head_logit.size(0						), self.n_token)						)
       UpperCAmelCase	:							Dict						=			nn.functional.log_softmax(__A,		dim=1						)

       UpperCAmelCase	:							List[str]						=			[0] + self.cutoffs
       for i in range(len(__A						) - 1						):
         UpperCAmelCase			,						UpperCAmelCase	:							Optional[Any]						=			cutoff_values[i], cutoff_values[i + 1]

         if i == 0:
           UpperCAmelCase	:							Any						=			head_logprob[:, : self.cutoffs[0]]
         else:
           UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							int						=			weights[i], biases[i], self.out_projs[i]

           UpperCAmelCase	:							Tuple						=			self._compute_logit(__A,		__A,		__A,		__A						)
           UpperCAmelCase	:							List[Any]						=			nn.functional.log_softmax(__A,		dim=1						)

           UpperCAmelCase	:							Optional[int]						=			head_logprob[:, -i] + tail_logprob_i
           UpperCAmelCase	:							Optional[Any]						=			logprob_i

       return out





