

"""simple docstring"""




from typing import Dict, List, Optional, Union

import numpy as np

from .feature_extraction_utils import BatchFeature, FeatureExtractionMixin
from .utils import PaddingStrategy, TensorType, is_tf_tensor, is_torch_tensor, logging, to_numpy


__UpperCAmelCase =				logging.get_logger(__name__)







class      __lowercase			(    __lowerCamelCase					):

  def __init__(  self		:			Optional[int]					,A		:			int					,A		:			int					,A		:			float					,**A		:			List[str]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					feature_size
   UpperCAmelCase__					:  int					   =					sampling_rate
   UpperCAmelCase__					:  List[Any]					   =					padding_value

   UpperCAmelCase__					:  int					   =					kwargs.pop("""padding_side"""					,"""right"""    )
   UpperCAmelCase__					:  int					   =					kwargs.pop("""return_attention_mask"""					,A    )

   super().__init__(**A    )

  def    __lowercase       (  self		:			Union[str, Any]					,A		:			Union[
          BatchFeature,
          List[BatchFeature],
          Dict[str, BatchFeature],
          Dict[str, List[BatchFeature]],
          List[Dict[str, BatchFeature]],
      ]					,A		:			Union[bool, str, PaddingStrategy] = True					,A		:			Optional[int] = None					,A		:			bool = False					,A		:			Optional[int] = None					,A		:			Optional[bool] = None					,A		:			Optional[Union[str, TensorType]] = None					,):






   '''simple docstring'''
   # If we have a list of dicts, let's convert it in a dict of lists
   # We do this to allow using this method as a collate_fn function in PyTorch Dataloader
   if isinstance(A					,(list, tuple)    ) and isinstance(processed_features[0]					,(dict, BatchFeature)    ):
    UpperCAmelCase__					:  Tuple					   =					{
        key: [example[key] for example in processed_features] for key in processed_features[0].keys()
    }

   # The model's main input name, usually `input_values`, has be passed for padding
   if self.model_input_names[0] not in processed_features:
    raise ValueError(
        """You should supply an instance of `transformers.BatchFeature` or list of `transformers.BatchFeature`"""
        f" to this method that includes {self.model_input_names[0]}, but you provided"
        f" {list(processed_features.keys()    )}"    )

   UpperCAmelCase__					:  str					   =					processed_features[self.model_input_names[0]]
   UpperCAmelCase__					:  Any					   =					(
       return_attention_mask if return_attention_mask is not None else self.return_attention_mask
   )

   if len(A    ) == 0:
    if return_attention_mask:
     UpperCAmelCase__					:  str					   =					[]
    return processed_features

   # If we have PyTorch/TF tensors or lists as inputs, we cast them as Numpy arrays
   # and rebuild them afterwards if no return_tensors is specified
   # Note that we lose the specific device the tensor may be on for PyTorch

   UpperCAmelCase__					:  int					   =					required_input[0]
   if isinstance(A					,(list, tuple)    ):
    # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element.
    UpperCAmelCase__					:  int					   =					0
    while len(required_input[index]    ) == 0:
     index += 1
    if index < len(A    ):
     UpperCAmelCase__					:  Union[str, Any]					   =					required_input[index][0]

   if return_tensors is None:
    if is_tf_tensor(A    ):
     UpperCAmelCase__					:  List[Any]					   =					"""tf"""
    elif is_torch_tensor(A    ):
     UpperCAmelCase__					:  Optional[int]					   =					"""pt"""
    elif isinstance(A					,(int, float, list, tuple, np.ndarray)    ):
     UpperCAmelCase__					:  List[Any]					   =					"""np"""
    else:
     raise ValueError(
         f"type of {first_element} unknown: {type(A    )}. "
         """Should be one of a python, numpy, pytorch or tensorflow object."""    )

   for key, value in processed_features.items():
    if isinstance(value[0]					,(int, float)    ):
     UpperCAmelCase__					:  Tuple					   =					to_numpy(A    )
    else:
     UpperCAmelCase__					:  str					   =					[to_numpy(A    ) for v in value]

        # Convert padding_strategy in PaddingStrategy
   UpperCAmelCase__					:  Any					   =					self._get_padding_strategies(padding=A					,max_length=A    )

   UpperCAmelCase__					:  Union[str, Any]					   =					processed_features[self.model_input_names[0]]

   UpperCAmelCase__					:  List[Any]					   =					len(A    )
   if not all(len(A    ) == batch_size for v in processed_features.values()    ):
    raise ValueError("""Some items in the output dictionary have a different batch size than others."""    )

   UpperCAmelCase__					:  List[str]					   =					[]
   for i in range(A    ):
    UpperCAmelCase__					:  Union[str, Any]					   =					{k: v[i] for k, v in processed_features.items()}
    # truncation
    UpperCAmelCase__					:  List[str]					   =					self._truncate(
        A					,max_length=A					,pad_to_multiple_of=A					,truncation=A					,)
    truncated_inputs.append(A    )

   if padding_strategy == PaddingStrategy.LONGEST:
    # make sure that `max_length` cannot be longer than the longest truncated length
    UpperCAmelCase__					:  Any					   =					max(len(input_slice[self.model_input_names[0]]    ) for input_slice in truncated_inputs    )
    UpperCAmelCase__					:  Dict					   =					PaddingStrategy.MAX_LENGTH

   UpperCAmelCase__					:  Optional[int]					   =					{}
   for i in range(A    ):
    # padding
    UpperCAmelCase__					:  List[str]					   =					self._pad(
        truncated_inputs[i]					,max_length=A					,padding_strategy=A					,pad_to_multiple_of=A					,return_attention_mask=A					,)

    for key, value in outputs.items():
     if key not in batch_outputs:
      UpperCAmelCase__					:  List[str]					   =					[]
     if value.dtype is np.dtype(np.floataa    ):
      UpperCAmelCase__					:  Optional[Any]					   =					value.astype(np.floataa    )
     batch_outputs[key].append(A    )

   return BatchFeature(A					,tensor_type=A    )

  def    __lowercase       (  self		:			List[Any]					,A		:			Union[Dict[str, np.ndarray], BatchFeature]					,A		:			Optional[int] = None					,A		:			PaddingStrategy = PaddingStrategy.DO_NOT_PAD					,A		:			Optional[int] = None					,A		:			Optional[bool] = None					,):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					processed_features[self.model_input_names[0]]

   if padding_strategy == PaddingStrategy.LONGEST:
    UpperCAmelCase__					:  Optional[int]					   =					len(A    )

   if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
    UpperCAmelCase__					:  str					   =					((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

   UpperCAmelCase__					:  Union[str, Any]					   =					padding_strategy != PaddingStrategy.DO_NOT_PAD and len(A    ) < max_length

   if return_attention_mask and "attention_mask" not in processed_features:
    UpperCAmelCase__					:  Tuple					   =					np.ones(len(A    )					,dtype=np.intaa    )

   if needs_to_be_padded:
    UpperCAmelCase__					:  Dict					   =					max_length - len(A    )
    if self.padding_side == "right":
     if return_attention_mask:
      UpperCAmelCase__					:  Dict					   =					np.pad(
          processed_features["""attention_mask"""]					,(0, difference)    )
     UpperCAmelCase__					:  List[str]					   =					((0, difference), (0, 0)) if self.feature_size > 1 else (0, difference)
     UpperCAmelCase__					:  Optional[int]					   =					np.pad(
         A					,A					,"""constant"""					,constant_values=self.padding_value    )
    elif self.padding_side == "left":
     if return_attention_mask:
      UpperCAmelCase__					:  Optional[Any]					   =					np.pad(
          processed_features["""attention_mask"""]					,(difference, 0)    )
     UpperCAmelCase__					:  Any					   =					((difference, 0), (0, 0)) if self.feature_size > 1 else (difference, 0)
     UpperCAmelCase__					:  List[str]					   =					np.pad(
         A					,A					,"""constant"""					,constant_values=self.padding_value    )
    else:
     raise ValueError("""Invalid padding strategy:""" + str(self.padding_side    )    )

   return processed_features

  def    __lowercase       (  self		:			List[str]					,A		:			Union[Dict[str, np.ndarray], BatchFeature]					,A		:			Optional[int] = None					,A		:			Optional[int] = None					,A		:			Optional[bool] = None					,):






   '''simple docstring'''
   if not truncation:
    return processed_features
   elif truncation and max_length is None:
    raise ValueError("""When setting ``truncation=True``, make sure that ``max_length`` is defined."""    )

   UpperCAmelCase__					:  List[Any]					   =					processed_features[self.model_input_names[0]]

   # find `max_length` that fits `pad_to_multiple_of`
   if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
    UpperCAmelCase__					:  Union[str, Any]					   =					((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

   UpperCAmelCase__					:  str					   =					len(A    ) > max_length

   if needs_to_be_truncated:
    UpperCAmelCase__					:  Optional[int]					   =					processed_features[self.model_input_names[0]][:max_length]
    if "attention_mask" in processed_features:
     UpperCAmelCase__					:  Dict					   =					processed_features["""attention_mask"""][:max_length]

   return processed_features




  def    __lowercase       (  self		:			Optional[Any]					,A		:			Dict=False					,A		:			Any=None    ):






   '''simple docstring'''

   # Get padding strategy
   if padding is not False:
    if padding is True:
     UpperCAmelCase__					:  List[str]					   =					PaddingStrategy.LONGEST  # Default to pad to the longest sequence in the batch
    elif not isinstance(A					,A    ):
     UpperCAmelCase__					:  int					   =					PaddingStrategy(A    )
    elif isinstance(A					,A    ):
     UpperCAmelCase__					:  Union[str, Any]					   =					padding
   else:
    UpperCAmelCase__					:  Tuple					   =					PaddingStrategy.DO_NOT_PAD

   # Set max length if needed
   if max_length is None:
    if padding_strategy == PaddingStrategy.MAX_LENGTH:
     raise ValueError(
         f"When setting ``padding={PaddingStrategy.MAX_LENGTH}``, make sure that max_length is defined"    )

        # Test if we have a padding value
   if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.padding_value is None):
    raise ValueError(
        """Asking to pad but the feature_extractor does not have a padding value. Please select a value to use"""
        """ as `padding_value`. For example: `feature_extractor.padding_value = 0.0`."""    )

   return padding_strategy








"""simple docstring"""




from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


__UpperCAmelCase =				logging.get_logger(__name__)


__UpperCAmelCase =				TypeVar('DatasetType', Dataset, IterableDataset)
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = "first_exhausted"							,    ):
 '''simple docstring'''





 from .arrow_dataset import Dataset
 from .iterable_dataset import IterableDataset

 if not datasets:
  raise ValueError("""Unable to interleave an empty list of datasets."""      )
 for i, dataset in enumerate(__UpperCamelCase      ):
  if not isinstance(__UpperCamelCase							,    (Dataset, IterableDataset)      ):
   if isinstance(__UpperCamelCase							,    (DatasetDict, IterableDatasetDict)      ):
    if not dataset:
     raise ValueError(
         F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
         """is an empty dataset dictionary."""      )
    raise ValueError(
        F"Dataset at position {i} has at least one split: {list(__UpperCamelCase      )}\n"
        F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase      )      )}']"      )
   raise ValueError(
       F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase      ).__name__}."      )
  if i == 0:
   UpperCAmelCase__   ,      UpperCAmelCase__					:  List[str]					   =					(
       (Dataset, IterableDataset) if isinstance(__UpperCamelCase							,    __UpperCamelCase      ) else (IterableDataset, Dataset)
   )
  elif not isinstance(__UpperCamelCase							,    __UpperCamelCase      ):
   raise ValueError(
       F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."      )
 if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
  raise ValueError(F"{stopping_strategy} is not supported. Please enter a valid stopping_strategy."      )
 if dataset_type is Dataset:
  return _interleave_map_style_datasets(
      __UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    stopping_strategy=__UpperCamelCase      )
 else:
  return _interleave_iterable_datasets(
      __UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    stopping_strategy=__UpperCamelCase      )




def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = 0							,    ):
 '''simple docstring'''





 if not dsets:
  raise ValueError("""Unable to concatenate an empty list of datasets."""      )
 for i, dataset in enumerate(__UpperCamelCase      ):
  if not isinstance(__UpperCamelCase							,    (Dataset, IterableDataset)      ):
   if isinstance(__UpperCamelCase							,    (DatasetDict, IterableDatasetDict)      ):
    if not dataset:
     raise ValueError(
         F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
         """is an empty dataset dictionary."""      )
    raise ValueError(
        F"Dataset at position {i} has at least one split: {list(__UpperCamelCase      )}\n"
        F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase      )      )}']"      )
   raise ValueError(
       F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase      ).__name__}."      )
  if i == 0:
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Optional[int]					   =					(
       (Dataset, IterableDataset) if isinstance(__UpperCamelCase							,    __UpperCamelCase      ) else (IterableDataset, Dataset)
   )
  elif not isinstance(__UpperCamelCase							,    __UpperCamelCase      ):
   raise ValueError(
       F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."      )
 if dataset_type is Dataset:
  return _concatenate_map_style_datasets(__UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    axis=__UpperCamelCase      )
 else:
  return _concatenate_iterable_datasets(__UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    axis=__UpperCamelCase      )





