"""simple docstring"""







# Lint as: python3
# pylint: enable=line-too-long
# pylint: disable=g-import-not-at-top,g-bad-import-order,wrong-import-position

__SCREAMING_SNAKE_CASE								="2.13.1"

import platform

import pyarrow
from packaging import version


if version.parse(platform.python_version()) < version.parse("3.7"):
       raise ImportWarning(
           "To use `datasets`, Python>=3.7 is required, and the current version of Python doesn't match this condition."
       )

if version.parse(pyarrow.__version__).major < 8:
       raise ImportWarning(
           "To use `datasets`, the module `pyarrow>=8.0.0` is required, and the current version of `pyarrow` doesn't match this condition.\n"
           "If you are running this in a Google Colab, you should probably just restart the runtime to use the right version of `pyarrow`."
       )

del platform
del pyarrow
del version

from .arrow_dataset import Dataset
from .arrow_reader import ReadInstruction
from .builder import ArrowBasedBuilder, BeamBasedBuilder, BuilderConfig, DatasetBuilder, GeneratorBasedBuilder
from .combine import concatenate_datasets, interleave_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .download import *
from .features import *
from .fingerprint import disable_caching, enable_caching, is_caching_enabled, set_caching_enabled
from .info import DatasetInfo, MetricInfo
from .inspect import (
    get_dataset_config_info,
    get_dataset_config_names,
    get_dataset_infos,
    get_dataset_split_names,
    inspect_dataset,
    inspect_metric,
    list_datasets,
    list_metrics,
)
from .iterable_dataset import IterableDataset
from .load import load_dataset, load_dataset_builder, load_from_disk, load_metric
from .metric import Metric
from .splits import (
    NamedSplit,
    NamedSplitAll,
    Split,
    SplitBase,
    SplitDict,
    SplitGenerator,
    SplitInfo,
    SubSplitInfo,
    percent,
)
from .tasks import *
from .utils import *
from .utils import logging


# deprecated modules
from datasets import arrow_dataset as _arrow_dataset  # isort:skip
from datasets import utils as _utils  # isort:skip
from datasets.utils import download_manager as _deprecated_download_manager  # isort:skip

__SCREAMING_SNAKE_CASE								=concatenate_datasets
__SCREAMING_SNAKE_CASE								=DownloadConfig
__SCREAMING_SNAKE_CASE								=DownloadManager
__SCREAMING_SNAKE_CASE								=DownloadMode
__SCREAMING_SNAKE_CASE								=DownloadConfig
__SCREAMING_SNAKE_CASE								=DownloadMode
__SCREAMING_SNAKE_CASE								=DownloadManager

del _arrow_dataset, _utils, _deprecated_download_manager


"""simple docstring"""







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


__SCREAMING_SNAKE_CASE								=logging.get_logger(__name__)







class  UpperCamelCase	(    lowercase_						):
   lowercase							=						['input_features', 'is_longer']
   def __init__(    self  ,__UpperCamelCase=64  ,__UpperCamelCase=4_8000  ,__UpperCamelCase=480  ,__UpperCamelCase=10  ,__UpperCamelCase=1024  ,__UpperCamelCase=0.0  ,__UpperCamelCase=False  ,__UpperCamelCase = 0  ,__UpperCamelCase = 1_4000  ,__UpperCamelCase = None  ,__UpperCamelCase = "fusion"  ,__UpperCamelCase = "repeatpad"  ,**__UpperCamelCase  ,)			->   Optional[int]:


        '''simple docstring'''


        super().__init__(
            feature_size=__UpperCamelCase  ,sampling_rate=__UpperCamelCase  ,padding_value=__UpperCamelCase  ,return_attention_mask=__UpperCamelCase  ,**__UpperCamelCase  ,)
        lowercase_			:       Union[str, Any]									=			top_db
        lowercase_			:       Any									=			truncation
        lowercase_			:       str									=			padding
        lowercase_			:       Optional[Any]									=			fft_window_size
        lowercase_			:       List[Any]									=			(fft_window_size >> 1) + 1
        lowercase_			:       Any									=			hop_length
        lowercase_			:       List[Any]									=			max_length_s
        lowercase_			:       Any									=			max_length_s * sampling_rate
        lowercase_			:       Optional[int]									=			sampling_rate
        lowercase_			:       List[Any]									=			frequency_min
        lowercase_			:       str									=			frequency_max
        lowercase_			:       Any									=			mel_filter_bank(
            num_frequency_bins=self.nb_frequency_bins  ,num_mel_filters=__UpperCamelCase  ,min_frequency=__UpperCamelCase  ,max_frequency=__UpperCamelCase  ,sampling_rate=__UpperCamelCase  ,norm=__UpperCamelCase  ,mel_scale='htk'  ,)
        lowercase_			:       Optional[Any]									=			mel_filter_bank(
            num_frequency_bins=self.nb_frequency_bins  ,num_mel_filters=__UpperCamelCase  ,min_frequency=__UpperCamelCase  ,max_frequency=__UpperCamelCase  ,sampling_rate=__UpperCamelCase  ,norm='slaney'  ,mel_scale='slaney'  ,)
   def     _UpperCAmelCase     (    self    )			->   Dict[str, Any]:


        '''simple docstring'''


        lowercase_			:       int									=			copy.deepcopy(self.__dict__    )
        lowercase_			:       Any									=			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     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase = None    )			->   np.ndarray:


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			spectrogram(
            __UpperCamelCase  ,window_function(self.fft_window_size  ,'hann'    )  ,frame_length=self.fft_window_size  ,hop_length=self.hop_length  ,power=2.0  ,mel_filters=__UpperCamelCase  ,log_mel='dB'  ,)
        return log_mel_spectrogram.T
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   Any:


        '''simple docstring'''


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

        lowercase_			:       Tuple									=			mel[idx_front : idx_front + chunk_frames, :]
        lowercase_			:       str									=			mel[idx_middle : idx_middle + chunk_frames, :]
        lowercase_			:       Optional[Any]									=			mel[idx_back : idx_back + chunk_frames, :]

        lowercase_			:       Tuple									=			torch.tensor(mel[None, None, :]    )
        lowercase_			:       Optional[Any]									=			torch.nn.functional.interpolate(
            __UpperCamelCase  ,size=[chunk_frames, 64]  ,mode='bilinear'  ,align_corners=__UpperCamelCase    )
        lowercase_			:       str									=			mel_shrink[0][0].numpy()
        lowercase_			:       Optional[int]									=			np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]  ,axis=0    )
        return mel_fusion
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   np.array:


        '''simple docstring'''


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

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

             if truncation == "fusion":
                  lowercase_			:       List[str]									=			self._np_extract_fbank_features(__UpperCamelCase  ,self.mel_filters    )
                  lowercase_			:       List[str]									=			np.stack([input_mel, input_mel, input_mel, input_mel]  ,axis=0    )
             else:
                  lowercase_			:       Any									=			self._np_extract_fbank_features(__UpperCamelCase  ,self.mel_filters_slaney    )[None, :]

        return input_mel, longer







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


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			truncation if truncation is not None else self.truncation
        lowercase_			:       List[Any]									=			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.'    )

        lowercase_			:       Tuple									=			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_			:       List[str]									=			is_batched_numpy or (
            isinstance(__UpperCamelCase  ,(list, tuple)    ) and (isinstance(raw_speech[0]  ,(np.ndarray, tuple, list)    ))
        )

        if is_batched:
             lowercase_			:       str									=			[np.asarray(__UpperCamelCase  ,dtype=np.floataa    ) for speech in raw_speech]
        elif not is_batched and not isinstance(__UpperCamelCase  ,np.ndarray    ):
             lowercase_			:       Any									=			np.asarray(__UpperCamelCase  ,dtype=np.floataa    )
        elif isinstance(__UpperCamelCase  ,np.ndarray    ) and raw_speech.dtype is np.dtype(np.floataa    ):
             lowercase_			:       Dict									=			raw_speech.astype(np.floataa    )

        # always return batch
        if not is_batched:
             lowercase_			:       int									=			[np.asarray(__UpperCamelCase    )]

        # convert to mel spectrogram, truncate and pad if needed.
        lowercase_			:       List[str]									=			[
            self._get_input_mel(__UpperCamelCase  ,max_length if max_length else self.nb_max_samples  ,__UpperCamelCase  ,__UpperCamelCase    )
            for waveform in raw_speech
        ]

        lowercase_			:       int									=			[]
        lowercase_			:       int									=			[]
        for mel, longer in padded_inputs:
             input_mel.append(__UpperCamelCase    )
             is_longer.append(__UpperCamelCase    )

        if truncation == "fusion" and sum(__UpperCamelCase    ) == 0:
             # if no audio is longer than 10s, then randomly select one audio to be longer
             lowercase_			:       int									=			np.random.randint(0  ,len(__UpperCamelCase    )    )
             lowercase_			:       Tuple									=			True

        if isinstance(input_mel[0]  ,__UpperCamelCase    ):
             lowercase_			:       Union[str, Any]									=			[np.asarray(__UpperCamelCase  ,dtype=np.floataa    ) for feature in input_mel]

        # is_longer is a list of bool
        lowercase_			:       Any									=			[[longer] for longer in is_longer]

        lowercase_			:       List[Any]									=			{'input_features': input_mel, 'is_longer': is_longer}
        lowercase_			:       Union[str, Any]									=			BatchFeature(__UpperCamelCase    )

        if return_tensors is not None:
             lowercase_			:       Any									=			input_features.convert_to_tensors(__UpperCamelCase    )

        return input_features

