




import os
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen, xsplitext
from ..table import array_cast
from ..utils.py_utils import no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
 from .features import FeatureType


_lowerCamelCase			,					_lowerCamelCase			,					_lowerCamelCase      			=						False, False, False



@dataclass
class       a					:

     '''simple docstring'''





     lowerCAmelCase    :	Optional[int]       		=  None
     lowerCAmelCase    :	bool       		=  True
     lowerCAmelCase    :	bool       		=  True
     lowerCAmelCase    :	Optional[str]       		=  None
     # Automatically constructed
     lowerCAmelCase    :	ClassVar[str]       		=  "dict"
     lowerCAmelCase    :	ClassVar[Any]       		=  pa.struct({'bytes': pa.binary(), 'path': pa.string()}      )
     lowerCAmelCase    :	str       		=  field(default='Audio'     ,       init=_A     ,       repr=_A      )
     def __call__(     self       :      Any    ):
         return self.pa_type
     def   lowerCamelCase_       (     self       :      Tuple ,   __snake_case       :      Union[str, bytes, dict]    ):
         try:
             import soundfile as sf  # soundfile is a dependency of librosa, needed to decode audio files.
         except ImportError as err:
             raise ImportError('''To support encoding audio data, please install \'soundfile\'.'''    ) from err
         if isinstance(__snake_case ,   __snake_case    ):
             return {"bytes": None, "path": value}
         elif isinstance(__snake_case ,   __snake_case    ):
             return {"bytes": value, "path": None}
         elif "array" in value:
             # convert the audio array to wav bytes
             UpperCAmelCase_										=					BytesIO()
             sf.write(__snake_case ,   value['''array'''] ,   value['''sampling_rate'''] ,   format='''wav'''    )
             return {"bytes": buffer.getvalue(), "path": None}
         elif value.get('''path'''    ) is not None and os.path.isfile(value['''path''']    ):
             # we set "bytes": None to not duplicate the data if they're already available locally
             if value["path"].endswith('''pcm'''    ):
                 # "PCM" only has raw audio bytes
                 if value.get('''sampling_rate'''    ) is None:
                     # At least, If you want to convert "PCM-byte" to "WAV-byte", you have to know sampling rate
                     raise KeyError('''To use PCM files, please specify a \'sampling_rate\' in Audio object'''    )
                 if value.get('''bytes'''    ):
                     # If we already had PCM-byte, we don`t have to make "read file, make bytes" (just use it!)
                     UpperCAmelCase_										=					np.frombuffer(value['''bytes'''] ,   dtype=np.intaa    ).astype(np.floataa    ) / 3_27_67
                 else:
                     UpperCAmelCase_										=					np.memmap(value['''path'''] ,   dtype='''h''' ,   mode='''r'''    ).astype(np.floataa    ) / 3_27_67

                 UpperCAmelCase_										=					BytesIO(bytes()    )
                 sf.write(__snake_case ,   __snake_case ,   value['''sampling_rate'''] ,   format='''wav'''    )
                 return {"bytes": buffer.getvalue(), "path": None}
             else:
                 return {"bytes": None, "path": value.get('''path'''    )}
         elif value.get('''bytes'''    ) is not None or value.get('''path'''    ) is not None:
             # store the audio bytes, and path is used to infer the audio format using the file extension
             return {"bytes": value.get('''bytes'''    ), "path": value.get('''path'''    )}
         else:
             raise ValueError(
                 F'An audio sample should have one of \'path\' or \'bytes\' but they are missing or None in {value}.'    )
     def   lowerCamelCase_       (     self       :      List[Any] ,   __snake_case       :      dict ,   __snake_case       :      Optional[Dict[str, Union[str, bool, None]]] = None    ):
         if not self.decode:
             raise RuntimeError('''Decoding is disabled for this feature. Please use Audio(decode=True) instead.'''    )

         UpperCAmelCase_			,							UpperCAmelCase_										=					(value['''path'''], BytesIO(value['''bytes''']    )) if value['''bytes'''] is not None else (value['''path'''], None)
         if path is None and file is None:
             raise ValueError(F'An audio sample should have one of \'path\' or \'bytes\' but both are None in {value}.'    )

         try:
             import librosa
             import soundfile as sf
         except ImportError as err:
             raise ImportError('''To support decoding audio files, please install \'librosa\' and \'soundfile\'.'''    ) from err

         UpperCAmelCase_										=					xsplitext(__snake_case    )[1][1:].lower() if path is not None else None
         if not config.IS_OPUS_SUPPORTED and audio_format == "opus":
             raise RuntimeError(
                 '''Decoding \'opus\' files requires system library \'libsndfile\'>=1.0.31, '''
                 '''You can try to update `soundfile` python library: `pip install "soundfile>=0.12.1"`. '''    )
         elif not config.IS_MP3_SUPPORTED and audio_format == "mp3":
             raise RuntimeError(
                 '''Decoding \'mp3\' files requires system library \'libsndfile\'>=1.1.0, '''
                 '''You can try to update `soundfile` python library: `pip install "soundfile>=0.12.1"`. '''    )

         if file is None:
             UpperCAmelCase_										=					token_per_repo_id or {}
             UpperCAmelCase_										=					path.split('''::'''    )[-1]
             try:
                 UpperCAmelCase_										=					string_to_dict(__snake_case ,   config.HUB_DATASETS_URL    )['''repo_id''']
                 UpperCAmelCase_										=					token_per_repo_id[repo_id]
             except (ValueError, KeyError):
                 UpperCAmelCase_										=					None

             with xopen(__snake_case ,   '''rb''' ,   use_auth_token=__snake_case    ) as f:
                 UpperCAmelCase_			,							UpperCAmelCase_										=					sf.read(__snake_case    )

         else:
             UpperCAmelCase_			,							UpperCAmelCase_										=					sf.read(__snake_case    )

         UpperCAmelCase_										=					array.T
         if self.mono:
             UpperCAmelCase_										=					librosa.to_mono(__snake_case    )
         if self.sampling_rate and self.sampling_rate != sampling_rate:
             UpperCAmelCase_										=					librosa.resample(__snake_case ,   orig_sr=__snake_case ,   target_sr=self.sampling_rate    )
             UpperCAmelCase_										=					self.sampling_rate

         return {"path": path, "array": array, "sampling_rate": sampling_rate}
     def   lowerCamelCase_       (     self       :      Union[str, Any]    ):
         from .features import Value

         if self.decode:
             raise ValueError('''Cannot flatten a decoded Audio feature.'''    )
         return {
             "bytes": Value('''binary'''    ),
             "path": Value('''string'''    ),
         }
     def   lowerCamelCase_       (     self       :      Optional[int] ,   __snake_case       :      Union[pa.StringArray, pa.StructArray]    ):
         if pa.types.is_string(storage.type    ):
             UpperCAmelCase_										=					pa.array([None] * len(__snake_case    ) ,   type=pa.binary()    )
             UpperCAmelCase_										=					pa.StructArray.from_arrays([bytes_array, storage] ,   ['''bytes''', '''path'''] ,   mask=storage.is_null()    )
         elif pa.types.is_binary(storage.type    ):
             UpperCAmelCase_										=					pa.array([None] * len(__snake_case    ) ,   type=pa.string()    )
             UpperCAmelCase_										=					pa.StructArray.from_arrays([storage, path_array] ,   ['''bytes''', '''path'''] ,   mask=storage.is_null()    )
         elif pa.types.is_struct(storage.type    ) and storage.type.get_all_field_indices('''array'''    ):
             UpperCAmelCase_										=					pa.array([Audio().encode_example(__snake_case    ) if x is not None else None for x in storage.to_pylist()]    )
         elif pa.types.is_struct(storage.type    ):
             if storage.type.get_field_index('''bytes'''    ) >= 0:
                 UpperCAmelCase_										=					storage.field('''bytes'''    )
             else:
                 UpperCAmelCase_										=					pa.array([None] * len(__snake_case    ) ,   type=pa.binary()    )
             if storage.type.get_field_index('''path'''    ) >= 0:
                 UpperCAmelCase_										=					storage.field('''path'''    )
             else:
                 UpperCAmelCase_										=					pa.array([None] * len(__snake_case    ) ,   type=pa.string()    )
             UpperCAmelCase_										=					pa.StructArray.from_arrays([bytes_array, path_array] ,   ['''bytes''', '''path'''] ,   mask=storage.is_null()    )
         return array_cast(__snake_case ,   self.pa_type    )




     def   lowerCamelCase_       (     self       :      str ,   __snake_case       :      pa.StructArray    ):

         @no_op_if_value_is_null
         def path_to_bytes(__snake_case       :      Optional[int]    ):
             with xopen(__snake_case ,   '''rb'''    ) as f:
                 UpperCAmelCase_										=					f.read()
             return bytes_

         UpperCAmelCase_										=					pa.array(
             [
                 (path_to_bytes(x['''path''']    ) if x['''bytes'''] is None else x['''bytes''']) if x is not None else None
                 for x in storage.to_pylist()
             ] ,   type=pa.binary() ,   )
         UpperCAmelCase_										=					pa.array(
             [os.path.basename(__snake_case    ) if path is not None else None for path in storage.field('''path'''    ).to_pylist()] ,   type=pa.string() ,   )
         UpperCAmelCase_										=					pa.StructArray.from_arrays([bytes_array, path_array] ,   ['''bytes''', '''path'''] ,   mask=bytes_array.is_null()    )
         return array_cast(__snake_case ,   self.pa_type    )











import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map




def  SCREAMING_SNAKE_CASE     (	__UpperCamelCase      :						Union[dict, list, tuple, torch.Tensor]				)  ->							List[Tuple[int, ...]]:
    UpperCAmelCase_										=					[]
    if isinstance(__UpperCamelCase	,					__UpperCamelCase				):
        for v in tree.values():
            shapes.extend(_fetch_dims(__UpperCamelCase				)				)
    elif isinstance(__UpperCamelCase	,					(list, tuple)				):
        for t in tree:
            shapes.extend(_fetch_dims(__UpperCamelCase				)				)
    elif isinstance(__UpperCamelCase	,					torch.Tensor				):
        shapes.append(tree.shape				)
    else:
        raise ValueError('''Not supported'''				)

    return shapes




@torch.jit.ignore
def  SCREAMING_SNAKE_CASE     (	__UpperCamelCase      :						int	,					__UpperCamelCase      :						Tuple[int, ...]				)  ->							Tuple[int, ...]:
    UpperCAmelCase_										=					[]
    for d in reversed(__UpperCamelCase				):
        idx.append(flat_idx % d				)
        UpperCAmelCase_										=					flat_idx // d

    return tuple(reversed(__UpperCamelCase				)				)




@torch.jit.ignore
def  SCREAMING_SNAKE_CASE     (	__UpperCamelCase      :						Sequence[int]	,					__UpperCamelCase      :						Sequence[int]	,					__UpperCamelCase      :						Sequence[int]	,					__UpperCamelCase      :						Optional[Sequence[bool]] = None	,					__UpperCamelCase      :						Optional[Sequence[bool]] = None	,					)  ->							List[Tuple[slice, ...]]:

    # start_edges and end_edges both indicate whether, starting from any given
    # dimension, the start/end index is at the top/bottom edge of the
    # corresponding tensor, modeled as a tree
    def reduce_edge_list(__UpperCamelCase      :						List[bool]				) -> None:
        UpperCAmelCase_										=					True
        for i in range(len(__UpperCamelCase				)				):
            UpperCAmelCase_										=					-1 * (i + 1)
            l[reversed_idx] &= tally
            UpperCAmelCase_										=					l[reversed_idx]

    if start_edges is None:
        UpperCAmelCase_										=					[s == 0 for s in start]
        reduce_edge_list(__UpperCamelCase				)
    if end_edges is None:
        UpperCAmelCase_										=					[e == (d - 1) for e, d in zip(__UpperCamelCase	,					__UpperCamelCase				)]
        reduce_edge_list(__UpperCamelCase				)

    # Base cases. Either start/end are empty and we're done, or the final,
    # one-dimensional tensor can be simply sliced
    if len(__UpperCamelCase				) == 0:
        return [()]
    elif len(__UpperCamelCase				) == 1:
        return [(slice(start[0]	,					end[0] + 1				),)]

    UpperCAmelCase_										=					[]
    UpperCAmelCase_										=					[]

    # Dimensions common to start and end can be selected directly
    for s, e in zip(__UpperCamelCase	,					__UpperCamelCase				):
        if s == e:
            path_list.append(slice(__UpperCamelCase	,					s + 1				)				)
        else:
            break

    UpperCAmelCase_										=					tuple(__UpperCamelCase				)
    UpperCAmelCase_										=					len(__UpperCamelCase				)

    # start == end, and we're done
    if divergence_idx == len(__UpperCamelCase				):
        return [path]

    def upper() -> Tuple[Tuple[slice, ...], ...]:
        assert start_edges is not None
        assert end_edges is not None

        UpperCAmelCase_										=					start[divergence_idx]
        return tuple(
            path + (slice(__UpperCamelCase	,					sdi + 1				),) + s
            for s in _get_minimal_slice_set(
                start[divergence_idx + 1 :]	,					[d - 1 for d in dims[divergence_idx + 1 :]]	,					dims[divergence_idx + 1 :]	,					start_edges=start_edges[divergence_idx + 1 :]	,					end_edges=[True for _ in end_edges[divergence_idx + 1 :]]	,					)				)

    def lower() -> Tuple[Tuple[slice, ...], ...]:
        assert start_edges is not None
        assert end_edges is not None

        UpperCAmelCase_										=					end[divergence_idx]
        return tuple(
            path + (slice(__UpperCamelCase	,					edi + 1				),) + s
            for s in _get_minimal_slice_set(
                [0 for _ in start[divergence_idx + 1 :]]	,					end[divergence_idx + 1 :]	,					dims[divergence_idx + 1 :]	,					start_edges=[True for _ in start_edges[divergence_idx + 1 :]]	,					end_edges=end_edges[divergence_idx + 1 :]	,					)				)

    # If both start and end are at the edges of the subtree rooted at
    # divergence_idx, we can just select the whole subtree at once
    if start_edges[divergence_idx] and end_edges[divergence_idx]:
        slices.append(path + (slice(start[divergence_idx]	,					end[divergence_idx] + 1				),)				)
    # If just start is at the edge, we can grab almost all of the subtree,
    # treating only the ragged bottom edge as an edge case
    elif start_edges[divergence_idx]:
        slices.append(path + (slice(start[divergence_idx]	,					end[divergence_idx]				),)				)
        slices.extend(lower()				)
    # Analogous to the previous case, but the top is ragged this time
    elif end_edges[divergence_idx]:
        slices.extend(upper()				)
        slices.append(path + (slice(start[divergence_idx] + 1	,					end[divergence_idx] + 1				),)				)
    # If both sides of the range are ragged, we need to handle both sides
    # separately. If there's contiguous meat in between them, we can index it
    # in one big chunk
    else:
        slices.extend(upper()				)
        UpperCAmelCase_										=					end[divergence_idx] - start[divergence_idx]
        if middle_ground > 1:
            slices.append(path + (slice(start[divergence_idx] + 1	,					end[divergence_idx]				),)				)
        slices.extend(lower()				)

    return slices




@torch.jit.ignore
def  SCREAMING_SNAKE_CASE     (	__UpperCamelCase      :						torch.Tensor	,					__UpperCamelCase      :						int	,					__UpperCamelCase      :						int	,					__UpperCamelCase      :						int				)  ->							torch.Tensor:
    UpperCAmelCase_										=					t.shape[:no_batch_dims]
    UpperCAmelCase_										=					list(_flat_idx_to_idx(__UpperCamelCase	,					__UpperCamelCase				)				)
    # _get_minimal_slice_set is inclusive
    UpperCAmelCase_										=					list(_flat_idx_to_idx(flat_end - 1	,					__UpperCamelCase				)				)

    # Get an ordered list of slices to perform
    UpperCAmelCase_										=					_get_minimal_slice_set(
        __UpperCamelCase	,					__UpperCamelCase	,					__UpperCamelCase	,					)

    UpperCAmelCase_										=					[t[s] for s in slices]

    return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]				) for s in sliced_tensors]				)




def  SCREAMING_SNAKE_CASE     (	__UpperCamelCase      :						Callable	,					__UpperCamelCase      :						Dict[str, Any]	,					__UpperCamelCase      :						int	,					__UpperCamelCase      :						int	,					__UpperCamelCase      :						bool = False	,					__UpperCamelCase      :						Any = None	,					__UpperCamelCase      :						bool = False	,					)  ->							Any:
    if not (len(__UpperCamelCase				) > 0):
        raise ValueError('''Must provide at least one input'''				)

    UpperCAmelCase_										=					[shape[:no_batch_dims] for shape in _fetch_dims(__UpperCamelCase				)]
    UpperCAmelCase_										=					tuple([max(__UpperCamelCase				) for s in zip(*__UpperCamelCase				)]				)

    def _prep_inputs(__UpperCamelCase      :						torch.Tensor				) -> torch.Tensor:
        if not low_mem:
            if not sum(t.shape[:no_batch_dims]				) == no_batch_dims:
                UpperCAmelCase_										=					t.expand(orig_batch_dims + t.shape[no_batch_dims:]				)
            UpperCAmelCase_										=					t.reshape(-1	,					*t.shape[no_batch_dims:]				)
        else:
            UpperCAmelCase_										=					t.expand(orig_batch_dims + t.shape[no_batch_dims:]				)
        return t

    UpperCAmelCase_										=					tensor_tree_map(_prep_inputs	,					__UpperCamelCase				)
    UpperCAmelCase_										=					None
    if _out is not None:
        UpperCAmelCase_										=					tensor_tree_map(lambda __UpperCamelCase				: t.view([-1] + list(t.shape[no_batch_dims:]				)				)	,					_out				)

    UpperCAmelCase_										=					1
    for d in orig_batch_dims:
        flat_batch_dim *= d

    UpperCAmelCase_										=					flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

    def _select_chunk(__UpperCamelCase      :						torch.Tensor				) -> torch.Tensor:
        return t[i : i + chunk_size] if t.shape[0] != 1 else t

    UpperCAmelCase_										=					0
    UpperCAmelCase_										=					prepped_outputs
    for _ in range(__UpperCamelCase				):
        # Chunk the input
        if not low_mem:
            UpperCAmelCase_										=					_select_chunk
        else:
            UpperCAmelCase_										=					partial(
                _chunk_slice	,					flat_start=__UpperCamelCase	,					flat_end=min(__UpperCamelCase	,					i + chunk_size				)	,					no_batch_dims=len(__UpperCamelCase				)	,					)

        UpperCAmelCase_										=					tensor_tree_map(__UpperCamelCase	,					__UpperCamelCase				)

        # Run the layer on the chunk
        UpperCAmelCase_										=					layer(**__UpperCamelCase				)

        # Allocate space for the output
        if out is None:
            UpperCAmelCase_										=					tensor_tree_map(lambda __UpperCamelCase				: t.new_zeros((flat_batch_dim,) + t.shape[1:]				)	,					__UpperCamelCase				)

        # Put the chunk in its pre-allocated space
        if isinstance(__UpperCamelCase	,					__UpperCamelCase				):

            def assign(__UpperCamelCase      :						dict	,					__UpperCamelCase      :						dict				) -> None:
                for k, v in da.items():
                    if isinstance(__UpperCamelCase	,					__UpperCamelCase				):
                        assign(__UpperCamelCase	,					da[k]				)
                    else:
                        if _add_into_out:
                            v[i : i + chunk_size] += da[k]
                        else:
                            UpperCAmelCase_										=					da[k]

            assign(__UpperCamelCase	,					__UpperCamelCase				)
        elif isinstance(__UpperCamelCase	,					__UpperCamelCase				):
            for xa, xa in zip(__UpperCamelCase	,					__UpperCamelCase				):
                if _add_into_out:
                    xa[i : i + chunk_size] += xa
                else:
                    UpperCAmelCase_										=					xa
        elif isinstance(__UpperCamelCase	,					torch.Tensor				):
            if _add_into_out:
                out[i : i + chunk_size] += output_chunk
            else:
                UpperCAmelCase_										=					output_chunk
        else:
            raise ValueError('''Not supported'''				)

        i += chunk_size

    UpperCAmelCase_										=					tensor_tree_map(lambda __UpperCamelCase				: t.view(orig_batch_dims + t.shape[1:]				)	,					__UpperCamelCase				)

    return out



class       a					:

     '''simple docstring'''





     def __init__(     self       :      List[Any] ,   __snake_case       :      int = 5_12 ,   ):
         UpperCAmelCase_										=					max_chunk_size
         UpperCAmelCase_										=					None
         UpperCAmelCase_										=					None
     def   lowerCamelCase_       (     self       :      List[Any] ,   __snake_case       :      Callable ,   __snake_case       :      tuple ,   __snake_case       :      int    ):
         logging.info('''Tuning chunk size...'''    )

         if min_chunk_size >= self.max_chunk_size:
             return min_chunk_size

         UpperCAmelCase_										=					[2**l for l in range(int(math.log(self.max_chunk_size ,   2    )    ) + 1    )]
         UpperCAmelCase_										=					[c for c in candidates if c > min_chunk_size]
         UpperCAmelCase_										=					[min_chunk_size] + candidates
         candidates[-1] += 4

         def test_chunk_size(__snake_case       :      int    ) -> bool:
             try:
                 with torch.no_grad():
                     fn(*__snake_case ,   chunk_size=__snake_case    )
                 return True
             except RuntimeError:
                 return False

         UpperCAmelCase_										=					0
         UpperCAmelCase_										=					len(__snake_case    ) - 1
         while i > min_viable_chunk_size_index:
             UpperCAmelCase_										=					test_chunk_size(candidates[i]    )
             if not viable:
                 UpperCAmelCase_										=					(min_viable_chunk_size_index + i) // 2
             else:
                 UpperCAmelCase_										=					i
                 UpperCAmelCase_										=					(i + len(__snake_case    ) - 1) // 2

         return candidates[min_viable_chunk_size_index]
     def   lowerCamelCase_       (     self       :      int ,   __snake_case       :      Iterable ,   __snake_case       :      Iterable    ):
         UpperCAmelCase_										=					True
         for aa, aa in zip(__snake_case ,   __snake_case    ):
             assert type(__snake_case    ) == type(__snake_case    )
             if isinstance(__snake_case ,   (list, tuple)    ):
                 consistent &= self._compare_arg_caches(__snake_case ,   __snake_case    )
             elif isinstance(__snake_case ,   __snake_case    ):
                 UpperCAmelCase_										=					[v for _, v in sorted(aa.items() ,   key=lambda __snake_case    : x[0]    )]
                 UpperCAmelCase_										=					[v for _, v in sorted(aa.items() ,   key=lambda __snake_case    : x[0]    )]
                 consistent &= self._compare_arg_caches(__snake_case ,   __snake_case    )
             else:
                 consistent &= aa == aa

         return consistent




     def   lowerCamelCase_       (     self       :      str ,   __snake_case       :      Callable ,   __snake_case       :      tuple ,   __snake_case       :      int ,   ):
         UpperCAmelCase_										=					True
         UpperCAmelCase_										=					tree_map(lambda __snake_case    : a.shape if isinstance(__snake_case ,   torch.Tensor    ) else a ,   __snake_case ,   __snake_case    )
         if self.cached_arg_data is not None:
             # If args have changed shape/value, we need to re-tune
             assert len(self.cached_arg_data    ) == len(__snake_case    )
             UpperCAmelCase_										=					self._compare_arg_caches(self.cached_arg_data ,   __snake_case    )
         else:
             # Otherwise, we can reuse the precomputed value
             UpperCAmelCase_										=					False

         if not consistent:
             UpperCAmelCase_										=					self._determine_favorable_chunk_size(
                 __snake_case ,   __snake_case ,   __snake_case ,   )
             UpperCAmelCase_										=					arg_data

         assert self.cached_chunk_size is not None

         return self.cached_chunk_size





