


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 PaddingStrategy, TensorType, logging


UpperCAmelCase__   	=							logging.get_logger(__name__)







class       lowerCamelCase__   (       lowerCAmelCase):
     SCREAMING_SNAKE_CASE__      			=			['''input_features''', '''attention_mask''']





     def __init__(self		,     UpperCAmelCase=8_0		,     UpperCAmelCase=1_6_0_0_0		,     UpperCAmelCase=8_0		,     UpperCAmelCase=0.0		,     UpperCAmelCase=True		,     UpperCAmelCase=True		,     UpperCAmelCase=True		,     **UpperCAmelCase		,     )     ->   Dict:
      super().__init__(feature_size=UpperCAmelCase		,     sampling_rate=UpperCAmelCase		,     padding_value=UpperCAmelCase		,     **UpperCAmelCase       )
      _lowercase							      =num_mel_bins
      _lowercase							      =do_ceptral_normalize
      _lowercase							      =normalize_means
      _lowercase							      =normalize_vars
      _lowercase							      =True





     def      __A					(self		,     UpperCAmelCase		,     )     ->   np.ndarray:
      _lowercase							      =waveform * (2**1_5)  # Kaldi compliance: 16-bit signed integers
      _lowercase							      =torch.from_numpy(UpperCAmelCase       ).unsqueeze(0       )
      _lowercase							      =ta_kaldi.fbank(UpperCAmelCase		,     num_mel_bins=self.num_mel_bins		,     sample_frequency=self.sampling_rate       )
      return features.numpy()





     @staticmethod
     def      __A					(UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase = True		,     UpperCAmelCase = True		,     UpperCAmelCase = 0.0		,     )     ->   np.ndarray:
      # make sure we normalize float32 arrays
      if normalize_means:
       _lowercase							      =x[:input_length].mean(axis=0       )
       _lowercase							      =np.subtract(UpperCAmelCase		,     UpperCAmelCase       )
      if normalize_vars:
       _lowercase							      =x[:input_length].std(axis=0       )
       _lowercase							      =np.divide(UpperCAmelCase		,     UpperCAmelCase       )

      if input_length < x.shape[0]:
       _lowercase							      =padding_value

      # make sure array is in float32
      _lowercase							      =x.astype(np.floataa       )

      return x





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase = None       )     ->   List[np.ndarray]:
      _lowercase							      =attention_mask.sum(-1       ) if attention_mask is not None else [x.shape[0] for x in input_features]
      return [
          self.utterance_cmvn(UpperCAmelCase		,     UpperCAmelCase		,     self.normalize_means		,     self.normalize_vars		,     self.padding_value       )
          for x, n in zip(UpperCAmelCase		,     UpperCAmelCase       )
      ]





     def __call__(self		,     UpperCAmelCase		,     UpperCAmelCase = False		,     UpperCAmelCase = None		,     UpperCAmelCase = False		,     UpperCAmelCase = None		,     UpperCAmelCase = None		,     UpperCAmelCase = None		,     UpperCAmelCase = None		,     **UpperCAmelCase		,     )     ->   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.'''       )

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

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

      # always return batch
      if not is_batched:
       _lowercase							      =[raw_speech]

      # extract fbank features
      _lowercase							      =[self._extract_fbank_features(UpperCAmelCase       ) for waveform in raw_speech]

      # convert into correct format for padding
      _lowercase							      =BatchFeature({'''input_features''': features}       )

      _lowercase							      =self.pad(
          UpperCAmelCase		,     padding=UpperCAmelCase		,     max_length=UpperCAmelCase		,     truncation=UpperCAmelCase		,     pad_to_multiple_of=UpperCAmelCase		,     return_attention_mask=UpperCAmelCase		,     **UpperCAmelCase		,     )

      # make sure list is in array format
      _lowercase							      =padded_inputs.get('''input_features'''       )
      if isinstance(input_features[0]		,     UpperCAmelCase       ):
       _lowercase							      =[np.asarray(UpperCAmelCase		,     dtype=np.floataa       ) for feature in input_features]

      _lowercase							      =padded_inputs.get('''attention_mask'''       )
      if attention_mask is not None:
       _lowercase							      =[np.asarray(UpperCAmelCase		,     dtype=np.intaa       ) for array in attention_mask]

      # Utterance-level cepstral mean and variance normalization
      if self.do_ceptral_normalize:
       _lowercase							      =(
           np.array(UpperCAmelCase		,     dtype=np.intaa       )
           if self._get_padding_strategies(UpperCAmelCase		,     max_length=UpperCAmelCase       ) is not PaddingStrategy.DO_NOT_PAD
           else None
       )
       _lowercase							      =self.normalize(
           padded_inputs['''input_features''']		,     attention_mask=UpperCAmelCase       )

      if return_tensors is not None:
       _lowercase							      =padded_inputs.convert_to_tensors(UpperCAmelCase       )

      return padded_inputs







# flake8: noqa
# Lint as: python3

from typing import Dict, List, Optional, Type

from .. import config
from ..utils import logging
from .formatting import (
    ArrowFormatter,
    CustomFormatter,
    Formatter,
    PandasFormatter,
    PythonFormatter,
    TensorFormatter,
    format_table,
    query_table,
)
from .np_formatter import NumpyFormatter


__UpperCamelCase					:       List[str]    					=      logging.get_logger(__name__)

__UpperCamelCase					:       Dict[Optional[str], Type[Formatter]]    					=      {}
__UpperCamelCase					:       Dict[Optional[str], str]    					=      {}
__UpperCamelCase					:       Dict[Optional[str], Exception]    					=      {}





def 		__A		(	__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase = None							,			)      ->       Optional[int]:
	a										=							aliases if aliases is not None else []
	if format_type in _FORMAT_TYPES:
		logger.warning(
		    f'Overwriting format type \'{format_type}\' ({_FORMAT_TYPES[format_type].__name__} -> {formatter_cls.__name__})'    )
	a										=							formatter_cls
	for alias in set(aliases + [format_type]    ):
		if alias in _FORMAT_TYPES_ALIASES:
			logger.warning(
			    f'Overwriting format type alias \'{alias}\' ({_FORMAT_TYPES_ALIASES[alias]} -> {format_type})'    )
		a										=							format_type





def 		__A		(	__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase = None    )      ->       List[str]:
	a										=							aliases if aliases is not None else []
	for alias in set(aliases + [format_type]    ):
		a										=							unavailable_error


# Here we define all the available formatting functions that can be used by `Dataset.set_format`
_register_formatter(PythonFormatter, None, aliases=["python"])
_register_formatter(ArrowFormatter, "arrow", aliases=["pa", "pyarrow"])
_register_formatter(NumpyFormatter, "numpy", aliases=["np"])
_register_formatter(PandasFormatter, "pandas", aliases=["pd"])
_register_formatter(CustomFormatter, "custom")

if config.TORCH_AVAILABLE:
		from .torch_formatter import TorchFormatter

		_register_formatter(TorchFormatter, "torch", aliases=["pt", "pytorch"])
else:
		__UpperCamelCase					:       str    					=      ValueError("PyTorch needs to be installed to be able to return PyTorch tensors.")
		_register_unavailable_formatter(_torch_error, "torch", aliases=["pt", "pytorch"])

if config.TF_AVAILABLE:
		from .tf_formatter import TFFormatter

		_register_formatter(TFFormatter, "tensorflow", aliases=["tf"])
else:
		__UpperCamelCase					:       List[str]    					=      ValueError("Tensorflow needs to be installed to be able to return Tensorflow tensors.")
		_register_unavailable_formatter(_tf_error, "tensorflow", aliases=["tf"])

if config.JAX_AVAILABLE:
		from .jax_formatter import JaxFormatter

		_register_formatter(JaxFormatter, "jax", aliases=[])
else:
		__UpperCamelCase					:       List[str]    					=      ValueError("JAX needs to be installed to be able to return JAX arrays.")
		_register_unavailable_formatter(_jax_error, "jax", aliases=[])





def 		__A		(	__lowerCamelCase    )      ->       Optional[str]:
	if format_type in _FORMAT_TYPES_ALIASES:
		return _FORMAT_TYPES_ALIASES[format_type]
	else:
		return format_type





def 		__A		(	__lowerCamelCase							,			**__lowerCamelCase    )      ->       Formatter:
	a										=							get_format_type_from_alias(__lowerCamelCase    )
	if format_type in _FORMAT_TYPES:
		return _FORMAT_TYPES[format_type](**__lowerCamelCase    )
	if format_type in _FORMAT_TYPES_ALIASES_UNAVAILABLE:
		raise _FORMAT_TYPES_ALIASES_UNAVAILABLE[format_type]
	else:
		raise ValueError(
		    f'Return type should be None or selected in {list(type for type in _FORMAT_TYPES.keys() if type != None    )}, but got \'{format_type}\''    )

