

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


snake_case    :						Dict  =						logging.get_logger(__name__)




class 							lowerCAmelCase__  (	UpperCamelCase						):
					__A     :    List[Any]           =						['input_features', 'attention_mask']







					def __init__(       self      :     Optional[Any]			,		_A      :     Optional[int]=80			,		_A      :     Tuple=1_6000			,		_A      :     Union[str, Any]=80			,		_A      :     Dict=0.0			,		_A      :     Union[str, Any]=True			,		_A      :     int=True			,		_A      :     Dict=True			,		**_A      :     int			,		):
											super().__init__(feature_size=_A			,		sampling_rate=_A			,		padding_value=_A			,		**_A)
											A__	:							Dict						    = num_mel_bins
											A__	:							str						    = do_ceptral_normalize
											A__	:							Union[str, Any]						    = normalize_means
											A__	:							Optional[Any]						    = normalize_vars
											A__	:							int						    = True







					def  _lowercase			(       self      :     Tuple			,		_A      :     np.ndarray			,		):
											A__	:							str						    = waveform * (2**15)  # Kaldi compliance: 16-bit signed integers
											A__	:							Any						    = torch.from_numpy(_A).unsqueeze(0)
											A__	:							List[Any]						    = ta_kaldi.fbank(_A			,		num_mel_bins=self.num_mel_bins			,		sample_frequency=self.sampling_rate)
											return features.numpy()







					@staticmethod
					def  _lowercase			(       _A      :     np.ndarray			,		_A      :     int			,		_A      :     Optional[bool] = True			,		_A      :     Optional[bool] = True			,		_A      :     float = 0.0			,		):
											# make sure we normalize float32 arrays
											if normalize_means:
																	A__	:							Optional[Any]						    = x[:input_length].mean(axis=0)
																	A__	:							Tuple						    = np.subtract(_A			,		_A)
											if normalize_vars:
																	A__	:							Optional[Any]						    = x[:input_length].std(axis=0)
																	A__	:							Tuple						    = np.divide(_A			,		_A)

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

											# make sure array is in float32
											A__	:							Union[str, Any]						    = x.astype(np.floataa)

											return x







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







					def __call__(       self      :     List[str]			,		_A      :     Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]			,		_A      :     Union[bool, str, PaddingStrategy] = False			,		_A      :     Optional[int] = None			,		_A      :     bool = False			,		_A      :     Optional[int] = None			,		_A      :     Optional[Union[str, TensorType]] = None			,		_A      :     Optional[int] = None			,		_A      :     Optional[bool] = None			,		**_A      :     List[Any]			,		):

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

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

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

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

											# extract fbank features
											A__	:							Tuple						    = [self._extract_fbank_features(_A) for waveform in raw_speech]

											# convert into correct format for padding
											A__	:							Union[str, Any]						    = BatchFeature({"input_features": features})

											A__	:							Any						    = self.pad(
											    _A			,		padding=_A			,		max_length=_A			,		truncation=_A			,		pad_to_multiple_of=_A			,		return_attention_mask=_A			,		**_A			,		)

											# make sure list is in array format
											A__	:							Union[str, Any]						    = padded_inputs.get("input_features")
											if isinstance(input_features[0]			,		_A):
																	A__	:							Union[str, Any]						    = [np.asarray(_A			,		dtype=np.floataa) for feature in input_features]

											A__	:							Any						    = padded_inputs.get("attention_mask")
											if attention_mask is not None:
																	A__	:							Dict						    = [np.asarray(_A			,		dtype=np.intaa) for array in attention_mask]

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

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

											return padded_inputs




import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader




class 							lowerCAmelCase__  (	UpperCamelCase						):







					def __init__(       self      :     Dict			,		_A      :     NestedDataStructureLike[PathLike]			,		_A      :     Optional[NamedSplit] = None			,		_A      :     Optional[Features] = None			,		_A      :     str = None			,		_A      :     bool = False			,		_A      :     bool = False			,		_A      :     Optional[str] = None			,		_A      :     Optional[int] = None			,		**_A      :     Dict			,		):
											super().__init__(
											    _A			,		split=_A			,		features=_A			,		cache_dir=_A			,		keep_in_memory=_A			,		streaming=_A			,		num_proc=_A			,		**_A			,		)
											A__	:							Optional[int]						    = field
											A__	:							Dict						    = path_or_paths if isinstance(_A			,		_A) else {self.split: path_or_paths}
											A__	:							int						    = Json(
											    cache_dir=_A			,		data_files=_A			,		features=_A			,		field=_A			,		**_A			,		)







					def  _lowercase			(       self      :     Dict):
											# Build iterable dataset
											if self.streaming:
																	A__	:							Union[str, Any]						    = self.builder.as_streaming_dataset(split=self.split)
											# Build regular (map-style) dataset
											else:
																	A__	:							Optional[Any]						    = None
																	A__	:							Any						    = None
																	A__	:							Tuple						    = None
																	A__	:							Tuple						    = None

																	self.builder.download_and_prepare(
																	    download_config=_A			,		download_mode=_A			,		verification_mode=_A			,		base_path=_A			,		num_proc=self.num_proc			,		)
																	A__	:							Any						    = self.builder.as_dataset(
																	    split=self.split			,		verification_mode=_A			,		in_memory=self.keep_in_memory)
											return dataset




class 							lowerCAmelCase__  :







					def __init__(       self      :     Union[str, Any]			,		_A      :     Dataset			,		_A      :     Union[PathLike, BinaryIO]			,		_A      :     Optional[int] = None			,		_A      :     Optional[int] = None			,		**_A      :     Any			,		):
											if num_proc is not None and num_proc <= 0:
																	raise ValueError(F'num_proc {num_proc} must be an integer > 0.')

											A__	:							Tuple						    = dataset
											A__	:							Tuple						    = path_or_buf
											A__	:							int						    = batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
											A__	:							Tuple						    = num_proc
											A__	:							Any						    = "utf-8"
											A__	:							Optional[int]						    = to_json_kwargs







					def  _lowercase			(       self      :     Union[str, Any]):
											A__	:							Any						    = self.to_json_kwargs.pop("path_or_buf"			,		_A)
											A__	:							Union[str, Any]						    = self.to_json_kwargs.pop("orient"			,		"records")
											A__	:							List[Any]						    = self.to_json_kwargs.pop("lines"			,		True if orient == "records" else False)
											A__	:							List[Any]						    = self.to_json_kwargs.pop("index"			,		False if orient in ["split", "table"] else True)
											A__	:							int						    = self.to_json_kwargs.pop("compression"			,		_A)

											if compression not in [None, "infer", "gzip", "bz2", "xz"]:
																	raise NotImplementedError(F'`datasets` currently does not support {compression} compression')

											if isinstance(self.path_or_buf			,		(str, bytes, os.PathLike)):
																	with fsspec.open(self.path_or_buf			,		"wb"			,		compression=_A) as buffer:
																							A__	:							int						    = self._write(file_obj=_A			,		orient=_A			,		lines=_A			,		index=_A			,		**self.to_json_kwargs)
											else:
																	if compression:
																							raise NotImplementedError(
																							    F'The compression parameter is not supported when writing to a buffer, but compression={compression}'
																							    " was passed. Please provide a local path instead.")
																	A__	:							Optional[Any]						    = self._write(
																	    file_obj=self.path_or_buf			,		orient=_A			,		lines=_A			,		index=_A			,		**self.to_json_kwargs)
											return written







					def  _lowercase			(       self      :     List[Any]			,		_A      :     int):
											A__    ,       A__    ,       A__    ,       A__    ,       A__	:							Union[str, Any]						    = args

											A__	:							Dict						    = query_table(
											    table=self.dataset.data			,		key=slice(_A			,		offset + self.batch_size)			,		indices=self.dataset._indices			,		)
											A__	:							Dict						    = batch.to_pandas().to_json(
											    path_or_buf=_A			,		orient=_A			,		lines=_A			,		index=_A			,		**_A)
											if not json_str.endswith("\n"):
																	json_str += "\n"
											return json_str.encode(self.encoding)







					def  _lowercase			(       self      :     List[str]			,		_A      :     BinaryIO			,		_A      :     str			,		_A      :     List[Any]			,		_A      :     Optional[Any]			,		**_A      :     Optional[Any]			,		):
											A__	:							Any						    = 0

											if self.num_proc is None or self.num_proc == 1:
																	for offset in logging.tqdm(
																	    range(0			,		len(self.dataset)			,		self.batch_size)			,		unit="ba"			,		disable=not logging.is_progress_bar_enabled()			,		desc="Creating json from Arrow format"			,		):
																							A__	:							List[str]						    = self._batch_json((offset, orient, lines, index, to_json_kwargs))
																							written += file_obj.write(_A)
											else:
																	A__    ,       A__	:							Union[str, Any]						    = len(self.dataset), self.batch_size
																	with multiprocessing.Pool(self.num_proc) as pool:
																							for json_str in logging.tqdm(
																							    pool.imap(
																							        self._batch_json			,		[(offset, orient, lines, index, to_json_kwargs) for offset in range(0			,		_A			,		_A)]			,		)			,		total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size			,		unit="ba"			,		disable=not logging.is_progress_bar_enabled()			,		desc="Creating json from Arrow format"			,		):
																													written += file_obj.write(_A)

											return written