



from typing import List, Optional, Union

import numpy as np

from ....audio_utils import mel_filter_bank, optimal_fft_length, spectrogram, window_function
from ....feature_extraction_sequence_utils import SequenceFeatureExtractor
from ....feature_extraction_utils import BatchFeature
from ....file_utils import PaddingStrategy, TensorType
from ....utils import logging


__UpperCAmelCase			  =			logging.get_logger(__name__)



class 							lowerCAmelCase_						(    UpperCamelCase__				):
       UpperCAmelCase__							:   int  	=    ["input_features", "attention_mask"]
       def __init__(     self,     SCREAMING_SNAKE_CASE_=80,     SCREAMING_SNAKE_CASE_=1_6000,     SCREAMING_SNAKE_CASE_=0.0,     SCREAMING_SNAKE_CASE_=10,     SCREAMING_SNAKE_CASE_=25,     SCREAMING_SNAKE_CASE_="hamming_window",     SCREAMING_SNAKE_CASE_=3_2768.0,     SCREAMING_SNAKE_CASE_=0.97,     SCREAMING_SNAKE_CASE_=1.0,     SCREAMING_SNAKE_CASE_=True,     SCREAMING_SNAKE_CASE_=True,     SCREAMING_SNAKE_CASE_=False,     **SCREAMING_SNAKE_CASE_,     )	->						Optional[Any]:
          super().__init__(feature_size=lowercase_,     sampling_rate=lowercase_,     padding_value=lowercase_,     **lowercase_			)

          UpperCamelCase       :    Optional[int]		      =      feature_size
          UpperCamelCase       :    List[str]		      =      sampling_rate
          UpperCamelCase       :    int		      =      padding_value
          UpperCamelCase       :    List[str]		      =      hop_length
          UpperCamelCase       :    int		      =      win_length
          UpperCamelCase       :    List[str]		      =      frame_signal_scale
          UpperCamelCase       :    Optional[Any]		      =      preemphasis_coeff
          UpperCamelCase       :    int		      =      mel_floor
          UpperCamelCase       :    List[str]		      =      normalize_means
          UpperCamelCase       :    List[str]		      =      normalize_vars
          UpperCamelCase       :    List[str]		      =      win_function
          UpperCamelCase       :    Tuple		      =      return_attention_mask

          UpperCamelCase       :    Union[str, Any]		      =      win_length * sampling_rate // 1000
          UpperCamelCase       :    int		      =      hop_length * sampling_rate // 1000

          UpperCamelCase       :    Union[str, Any]		      =      optimal_fft_length(self.sample_size			)
          UpperCamelCase       :    Any		      =      (self.n_fft // 2) + 1
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_			)	->						np.ndarray:
          if self.win_function == "hamming_window":
             UpperCamelCase       :    Optional[int]		      =      window_function(window_length=self.sample_size,     name=self.win_function,     periodic=lowercase_			)
          else:
             UpperCamelCase       :    Optional[int]		      =      window_function(window_length=self.sample_size,     name=self.win_function			)

          UpperCamelCase       :    List[str]		      =      mel_filter_bank(
              num_frequency_bins=self.n_freqs,     num_mel_filters=self.feature_size,     min_frequency=0.0,     max_frequency=self.sampling_rate / 2.0,     sampling_rate=self.sampling_rate,     )

          UpperCamelCase       :    Optional[Any]		      =      spectrogram(
              one_waveform * self.frame_signal_scale,     window=lowercase_,     frame_length=self.sample_size,     hop_length=self.sample_stride,     fft_length=self.n_fft,     center=lowercase_,     preemphasis=self.preemphasis_coeff,     mel_filters=lowercase_,     mel_floor=self.mel_floor,     log_mel='log',     )
          return msfc_features.T
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						Dict:
          # make sure we normalize float32 arrays
          if self.normalize_means:
             UpperCamelCase       :    Optional[Any]		      =      x[:input_length].mean(axis=0			)
             UpperCamelCase       :    Any		      =      np.subtract(lowercase_,     lowercase_			)
          if self.normalize_vars:
             UpperCamelCase       :    Any		      =      x[:input_length].std(axis=0			)
             UpperCamelCase       :    int		      =      np.divide(lowercase_,     lowercase_			)

          if input_length < x.shape[0]:
             UpperCamelCase       :    Union[str, Any]		      =      padding_value

          # make sure array is in float32
          UpperCamelCase       :    Any		      =      x.astype(np.floataa			)

          return x
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_ = None			)	->						List[np.ndarray]:
          UpperCamelCase       :    List[str]		      =      attention_mask.sum(-1			) if attention_mask is not None else [x.shape[0] for x in input_features]
          return [self._normalize_one(lowercase_,     lowercase_,     self.padding_value			) for x, n in zip(lowercase_,     lowercase_			)]






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

          UpperCamelCase       :    Optional[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}"""			)
          UpperCamelCase       :    Optional[Any]		      =      is_batched_numpy or (
              isinstance(lowercase_,     (list, tuple)			) and (isinstance(raw_speech[0],     (np.ndarray, tuple, list)			))
          )

          if is_batched:
             UpperCamelCase       :    int		      =      [np.asarray(lowercase_,     dtype=np.floataa			) for speech in raw_speech]
          elif not is_batched and not isinstance(lowercase_,     np.ndarray			):
             UpperCamelCase       :    Tuple		      =      np.asarray(lowercase_,     dtype=np.floataa			)
          elif isinstance(lowercase_,     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       :    Union[str, Any]		      =      [raw_speech]

          # extract fbank features
          UpperCamelCase       :    List[str]		      =      [self._extract_mfsc_features(lowercase_			) for one_waveform in raw_speech]

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

          UpperCamelCase       :    str		      =      self.pad(
              lowercase_,     padding=lowercase_,     max_length=lowercase_,     truncation=lowercase_,     pad_to_multiple_of=lowercase_,     return_attention_mask=lowercase_,     **lowercase_,     )
          # make sure list is in array format
          UpperCamelCase       :    List[Any]		      =      padded_inputs.get('input_features'			)
          if isinstance(input_features[0],     lowercase_			):
             UpperCamelCase       :    Dict		      =      [np.asarray(lowercase_,     dtype=np.floataa			) for feature in input_features]

          UpperCamelCase       :    str		      =      padded_inputs.get('attention_mask'			)
          if attention_mask is not None:
             UpperCamelCase       :    Any		      =      [np.asarray(lowercase_,     dtype=np.intaa			) for array in attention_mask]

          if self.normalize_means or self.normalize_vars:
             UpperCamelCase       :    Any		      =      (
                 np.array(lowercase_,     dtype=np.intaa			)
                 if self._get_padding_strategies(lowercase_,     max_length=lowercase_			) is not PaddingStrategy.DO_NOT_PAD
                 and padding
                 else None
             )
             UpperCamelCase       :    List[str]		      =      self.normalize(
                 padded_inputs['input_features'],     attention_mask=lowercase_			)

          if return_tensors is not None:
             UpperCamelCase       :    Tuple		      =      padded_inputs.convert_to_tensors(lowercase_			)

          return padded_inputs










'''simple docstring'''

import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__lowerCAmelCase				:      int						      =5_0000
__lowerCAmelCase				:      Union[str, Any]						      =5000

__lowerCAmelCase      ,	__lowerCAmelCase				:      Tuple						      =os.path.split(__file__)
__lowerCAmelCase				:      int						      =os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   int			):
    for i in range(_lowerCamelCase			):
        A__							=     dataset[i]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[int]		,			_lowerCamelCase   :   Any			):
    for i in range(0		,			len(_lowerCamelCase			)		,			_lowerCamelCase			):
        A__							=     dataset[i : i + batch_size]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[Any]		,			_lowerCamelCase   :   List[str]			):
    with dataset.formatted_as(type=_lowerCamelCase			):
        for i in range(_lowerCamelCase			):
            A__							=     dataset[i]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[int]		,			_lowerCamelCase   :   int		,			_lowerCamelCase   :   Dict			):
    with dataset.formatted_as(type=_lowerCamelCase			):
        for i in range(0		,			_lowerCamelCase		,			_lowerCamelCase			):
            A__							=     dataset[i : i + batch_size]





def 		UpperCamelCase							(				):
    A__							=     {"num examples": SPEED_TEST_N_EXAMPLES}
    A__							=     [
        (read, {"length": SMALL_TEST}),
        (read, {"length": SPEED_TEST_N_EXAMPLES}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1_00}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10_00}),
        (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
        (read_formatted, {"type": "pandas", "length": SMALL_TEST}),
        (read_formatted, {"type": "torch", "length": SMALL_TEST}),
        (read_formatted, {"type": "tensorflow", "length": SMALL_TEST}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10_00}),
    ]

    A__							=     [
        (read, {"length": SMALL_TEST}),
        (read, {"length": SPEED_TEST_N_EXAMPLES}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1_00}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10_00}),
        (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10_00}),
    ]
    with tempfile.TemporaryDirectory() as tmp_dir:
        print("generating dataset"			)
        A__							=     datasets.Features(
            {"list": datasets.Sequence(datasets.Value("float32"			)			), "numbers": datasets.Value("float32"			)}			)
        A__							=     generate_example_dataset(
            os.path.join(_lowerCamelCase		,			"dataset.arrow"			)		,			_lowerCamelCase		,			num_examples=_lowerCamelCase		,			seq_shapes={"list": (1_00,)}		,			)
        print("first set of iterations"			)
        for func, kwargs in functions:
            print(func.__name__		,			str(_lowerCamelCase			)			)
            A__							=     func(_lowerCamelCase		,			**_lowerCamelCase			)

        print("shuffling dataset"			)
        A__							=     dataset.shuffle()
        print("Second set of iterations (after shuffling"			)
        for func, kwargs in functions_shuffled:
            print("shuffled "		,			func.__name__		,			str(_lowerCamelCase			)			)
            A__							=     func(
                _lowerCamelCase		,			**_lowerCamelCase			)

    with open(_lowerCamelCase		,			"wb"			) as f:
        f.write(json.dumps(_lowerCamelCase			).encode("utf-8"			)			)


if __name__ == "__main__":  # useful to run the profiler
  benchmark_iterating()



