


from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


_UpperCAmelCase      :	List[Any]  =logging.get_logger(__name__)







class 					snake_case__(				UpperCAmelCase__				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	Any       =				["""input_values""", """attention_mask"""]




						def __init__(					self ,    __lowercase=1 ,    __lowercase=1_6_0_0_0 ,    __lowercase=1_2_8 ,    __lowercase=1_0_2_4 ,    __lowercase=0.0 ,    __lowercase=True ,    __lowercase=-4.2_67_73_93 ,    __lowercase=4.5_68_99_74 ,    __lowercase=False ,    **__lowercase ,    )       ->     Any:
										super().__init__(feature_size=__lowercase ,    sampling_rate=__lowercase ,    padding_value=__lowercase ,    **__lowercase				)
										lowerCAmelCase_			:   Optional[Any]     					=					num_mel_bins
										lowerCAmelCase_			:   Tuple     					=					max_length
										lowerCAmelCase_			:   str     					=					do_normalize
										lowerCAmelCase_			:   List[Any]     					=					mean
										lowerCAmelCase_			:   List[Any]     					=					std
										lowerCAmelCase_			:   List[str]     					=					return_attention_mask




						def     lowercase_			(					self ,    __lowercase ,    __lowercase ,    )       ->     np.ndarray:
										lowerCAmelCase_			:   Union[str, Any]     					=					torch.from_numpy(__lowercase				).unsqueeze(0				)
										lowerCAmelCase_			:   List[str]     					=					ta_kaldi.fbank(
										    __lowercase ,    htk_compat=__lowercase ,    sample_frequency=self.sampling_rate ,    use_energy=__lowercase ,    window_type='''hanning''' ,    num_mel_bins=self.num_mel_bins ,    dither=0.0 ,    frame_shift=1_0 ,    )

										lowerCAmelCase_			:   Tuple     					=					fbank.shape[0]
										lowerCAmelCase_			:   Any     					=					max_length - n_frames

										# pad or truncate, depending on difference
										if difference > 0:
														lowerCAmelCase_			:   Dict     					=					torch.nn.ZeroPadad((0, 0, 0, difference)				)
														lowerCAmelCase_			:   Dict     					=					pad_module(__lowercase				)
										elif difference < 0:
														lowerCAmelCase_			:   Optional[Any]     					=					fbank[0:max_length, :]

										lowerCAmelCase_			:   Optional[Any]     					=					fbank.numpy()

										return fbank




						def     lowercase_			(					self ,    __lowercase				)       ->     np.ndarray:
										return (input_values - (self.mean)) / (self.std * 2)




						def __call__(					self ,    __lowercase ,    __lowercase = None ,    __lowercase = None ,    **__lowercase ,    )       ->     BatchFeature:

										if sampling_rate is not None:
														if sampling_rate != self.sampling_rate:
																		raise ValueError(
																		    f"""The model corresponding to this feature extractor: {self} was trained using a sampling rate of"""
																		    f""" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"""
																		    f""" {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_			:   Optional[int]     					=					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_			:   Tuple     					=					is_batched_numpy or (
										    isinstance(__lowercase ,    (list, tuple)				) and (isinstance(raw_speech[0] ,    (np.ndarray, tuple, list)				))
										)

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

										# always return batch
										if not is_batched:
														lowerCAmelCase_			:   Optional[Any]     					=					[raw_speech]

										# extract fbank features and pad/truncate to max_length
										lowerCAmelCase_			:   Any     					=					[self._extract_fbank_features(__lowercase ,    max_length=self.max_length				) for waveform in raw_speech]

										# convert into BatchFeature
										lowerCAmelCase_			:   List[Any]     					=					BatchFeature({'''input_values''': features}				)

										# make sure list is in array format
										lowerCAmelCase_			:   List[Any]     					=					padded_inputs.get('''input_values'''				)
										if isinstance(input_values[0] ,    __lowercase				):
														lowerCAmelCase_			:   List[str]     					=					[np.asarray(__lowercase ,    dtype=np.floataa				) for feature in input_values]

										# normalization
										if self.do_normalize:
														lowerCAmelCase_			:   Union[str, Any]     					=					[self.normalize(__lowercase				) for feature in input_values]

										if return_tensors is not None:
														lowerCAmelCase_			:   Tuple     					=					padded_inputs.convert_to_tensors(__lowercase				)

										return padded_inputs



import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
						import transformer_engine.pytorch as te



def       lowerCAmelCase			(		lowerCAmelCase_			, lowerCAmelCase_=True			, lowerCAmelCase_=True			, lowerCAmelCase_=True			)->			Union[str, Any]:
				if not is_fpa_available():
								raise ImportError('''Using `convert_model` requires transformer_engine to be installed.'''			)
				for name, module in model.named_children():
								if isinstance(lowerCAmelCase_			, nn.Linear			) and to_transformer_engine and _convert_linear:
												# Return early if the linear layer weights are not multiples of 16
												if any(p % 16 != 0 for p in module.weight.shape			):
																return
												lowerCAmelCase_			:   Optional[int]     					=					module.bias is not None
												lowerCAmelCase_			:   Any     					=					te.Linear(
												    module.in_features			, module.out_features			, bias=lowerCAmelCase_			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Dict     					=					module.weight.data.clone()
												if has_bias:
																lowerCAmelCase_			:   Any     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, nn.LayerNorm			) and to_transformer_engine and _convert_ln:
												lowerCAmelCase_			:   List[str]     					=					te.LayerNorm(module.normalized_shape[0]			, eps=module.eps			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Union[str, Any]     					=					module.weight.data.clone()
												lowerCAmelCase_			:   List[Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, te.Linear			) and not to_transformer_engine and _convert_linear:
												lowerCAmelCase_			:   Dict     					=					module.bias is not None
												lowerCAmelCase_			:   Union[str, Any]     					=					nn.Linear(
												    module.in_features			, module.out_features			, bias=lowerCAmelCase_			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Optional[Any]     					=					module.weight.data.clone()
												if has_bias:
																lowerCAmelCase_			:   Union[str, Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, te.LayerNorm			) and not to_transformer_engine and _convert_ln:
												lowerCAmelCase_			:   Union[str, Any]     					=					nn.LayerNorm(module.normalized_shape[0]			, eps=module.eps			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   List[str]     					=					module.weight.data.clone()
												lowerCAmelCase_			:   List[Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								else:
												convert_model(
												    lowerCAmelCase_			, to_transformer_engine=lowerCAmelCase_			, _convert_linear=_convert_linear			, _convert_ln=_convert_ln			, )




def       lowerCAmelCase			(		lowerCAmelCase_			)->			str:
				if not is_fpa_available():
								raise ImportError('''Using `has_transformer_engine_layers` requires transformer_engine to be installed.'''			)
				for m in model.modules():
								if isinstance(lowerCAmelCase_			, (te.LayerNorm, te.Linear)			):
												return True
				return False