import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
  from transformers import SpeechaTextFeatureExtractor

UpperCAmelCase_			    =							random.Random()




def       lowerCAmelCase_     (		lowercase:							Optional[Any]  ,					lowercase:							Any=1.0  ,					lowercase:							Optional[Any]=None  ,					lowercase:							int=None		)    ->				Any:







 '''simple docstring'''
 if rng is None:
  _UpperCamelCase:     Union[str, Any]							= global_rng

 _UpperCamelCase:     Tuple							= []
 for batch_idx in range(shape[0]		):
  values.append([]		)
  for _ in range(shape[1]		):
   values[-1].append(rng.random() * scale		)

 return values




@require_torch
@require_torchaudio
class  __magic_name__  (		unittest.TestCase    ):





   """simple docstring"""




   def __init__(	self					:  Dict		,  _lowercase					:  Dict		,  _lowercase					:  Optional[Any]=7		,  _lowercase					:  Union[str, Any]=400		,  _lowercase					:  Any=2_000		,  _lowercase					:  str=24		,  _lowercase					:  List[Any]=24		,  _lowercase					:  Dict=0.0		,  _lowercase					:  int=16_000		,  _lowercase					:  Tuple=True		,  _lowercase					:  List[str]=True		,  ):



    """simple docstring"""



    _UpperCamelCase:     int							= parent
    _UpperCamelCase:     int							= batch_size
    _UpperCamelCase:     Any							= min_seq_length
    _UpperCamelCase:     List[str]							= max_seq_length
    _UpperCamelCase:     List[Any]							= (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
    _UpperCamelCase:     int							= feature_size
    _UpperCamelCase:     List[str]							= num_mel_bins
    _UpperCamelCase:     int							= padding_value
    _UpperCamelCase:     Optional[int]							= sampling_rate
    _UpperCamelCase:     Dict							= return_attention_mask
    _UpperCamelCase:     Tuple							= do_normalize



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    return {
        "feature_size": self.feature_size,
        "num_mel_bins": self.num_mel_bins,
        "padding_value": self.padding_value,
        "sampling_rate": self.sampling_rate,
        "return_attention_mask": self.return_attention_mask,
        "do_normalize": self.do_normalize,
    }



   def   lowerCAmelCase						(	self					:  List[Any]		,  _lowercase					:  Optional[int]=False		,  _lowercase					:  str=False							):



    """simple docstring"""



    def _flatten(_lowercase					:  Dict							):
     return list(itertools.chain(*_lowercase							)							)

    if equal_length:
     _UpperCamelCase:     Any							= [floats_list((self.max_seq_length, self.feature_size)							) for _ in range(self.batch_size							)]
    else:
     # make sure that inputs increase in size
     _UpperCamelCase:     str							= [
         floats_list((x, self.feature_size)							)
         for x in range(self.min_seq_length		,  self.max_seq_length		,  self.seq_length_diff							)
     ]
    if numpify:
     _UpperCamelCase:     str							= [np.asarray(_lowercase							) for x in speech_inputs]
    return speech_inputs




@require_torch
@require_torchaudio
class  __magic_name__  (		__a      ,	unittest.TestCase    ):





   """simple docstring"""




   lowerCAmelCase      :    Optional[int]  			=   SpeechaTextFeatureExtractor if is_speech_available() else None



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    _UpperCamelCase:     Union[str, Any]							= SpeechaTextFeatureExtractionTester(self							)



   def   lowerCAmelCase						(	self					:  List[Any]		,  _lowercase					:  Any							):



    """simple docstring"""



    self.assertTrue(np.all(np.mean(_lowercase		,  axis=0							) < 1E-3							)							)
    self.assertTrue(np.all(np.abs(np.var(_lowercase		,  axis=0							) - 1							) < 1E-3							)							)



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    _UpperCamelCase:     List[str]							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
    # create three inputs of length 800, 1000, and 1200
    _UpperCamelCase:     Tuple							= [floats_list((1, x)							)[0] for x in range(800		,  1_400		,  200							)]
    _UpperCamelCase:     List[str]							= [np.asarray(_lowercase							) for speech_input in speech_inputs]

    # Test feature size
    _UpperCamelCase:     str							= feature_extractor(_lowercase		,  padding=_lowercase		,  return_tensors='''np'''							).input_features
    self.assertTrue(input_features.ndim == 3							)
    self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size							)

    # Test not batched input
    _UpperCamelCase:     Union[str, Any]							= feature_extractor(speech_inputs[0]		,  return_tensors='''np'''							).input_features
    _UpperCamelCase:     Dict							= feature_extractor(np_speech_inputs[0]		,  return_tensors='''np'''							).input_features
    self.assertTrue(np.allclose(_lowercase		,  _lowercase		,  atol=1E-3							)							)

    # Test batched
    _UpperCamelCase:     str							= feature_extractor(_lowercase		,  return_tensors='''np'''							).input_features
    _UpperCamelCase:     Any							= feature_extractor(_lowercase		,  return_tensors='''np'''							).input_features
    for enc_seq_a, enc_seq_a in zip(_lowercase		,  _lowercase							):
     self.assertTrue(np.allclose(_lowercase		,  _lowercase		,  atol=1E-3							)							)

    # Test 2-D numpy arrays are batched.
    _UpperCamelCase:     Tuple							= [floats_list((1, x)							)[0] for x in (800, 800, 800)]
    _UpperCamelCase:     Optional[int]							= np.asarray(_lowercase							)
    _UpperCamelCase:     Tuple							= feature_extractor(_lowercase		,  return_tensors='''np'''							).input_features
    _UpperCamelCase:     List[str]							= feature_extractor(_lowercase		,  return_tensors='''np'''							).input_features
    for enc_seq_a, enc_seq_a in zip(_lowercase		,  _lowercase							):
     self.assertTrue(np.allclose(_lowercase		,  _lowercase		,  atol=1E-3							)							)



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase:     Dict							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
    _UpperCamelCase:     List[str]							= [floats_list((1, x)							)[0] for x in range(800		,  1_400		,  200							)]

    _UpperCamelCase:     Dict							= ['''longest''', '''max_length''', '''do_not_pad''']
    _UpperCamelCase:     Tuple							= [None, 16, None]
    for max_length, padding in zip(_lowercase		,  _lowercase							):
     _UpperCamelCase:     List[Any]							= feature_extractor(
         _lowercase		,  padding=_lowercase		,  max_length=_lowercase		,  return_attention_mask=_lowercase							)
     _UpperCamelCase:     str							= inputs.input_features
     _UpperCamelCase:     List[Any]							= inputs.attention_mask
     _UpperCamelCase:     List[Any]							= [np.sum(_lowercase							) for x in attention_mask]

     self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
     self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
     self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)



   def   lowerCAmelCase						(	self					:  str							):



    """simple docstring"""



    _UpperCamelCase:     Any							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
    _UpperCamelCase:     Optional[Any]							= [floats_list((1, x)							)[0] for x in range(800		,  1_400		,  200							)]

    _UpperCamelCase:     Union[str, Any]							= ['''longest''', '''max_length''', '''do_not_pad''']
    _UpperCamelCase:     List[Any]							= [None, 16, None]
    for max_length, padding in zip(_lowercase		,  _lowercase							):
     _UpperCamelCase:     Tuple							= feature_extractor(
         _lowercase		,  max_length=_lowercase		,  padding=_lowercase		,  return_tensors='''np'''		,  return_attention_mask=_lowercase							)
     _UpperCamelCase:     Optional[int]							= inputs.input_features
     _UpperCamelCase:     Union[str, Any]							= inputs.attention_mask
     _UpperCamelCase:     Union[str, Any]							= [np.sum(_lowercase							) for x in attention_mask]

     self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
     self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6							)
     self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
     self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6							)
     self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)



   def   lowerCAmelCase						(	self					:  int							):



    """simple docstring"""



    _UpperCamelCase:     Dict							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
    _UpperCamelCase:     Union[str, Any]							= [floats_list((1, x)							)[0] for x in range(800		,  1_400		,  200							)]
    _UpperCamelCase:     List[Any]							= feature_extractor(
        _lowercase		,  padding='''max_length'''		,  max_length=4		,  truncation=_lowercase		,  return_tensors='''np'''		,  return_attention_mask=_lowercase		,  )
    _UpperCamelCase:     List[Any]							= inputs.input_features
    _UpperCamelCase:     Any							= inputs.attention_mask
    _UpperCamelCase:     Union[str, Any]							= np.sum(attention_mask == 1		,  axis=1							)

    self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
    self._check_zero_mean_unit_variance(input_features[1]							)
    self._check_zero_mean_unit_variance(input_features[2]							)



   def   lowerCAmelCase						(	self					:  int							):



    """simple docstring"""



    _UpperCamelCase:     List[Any]							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
    _UpperCamelCase:     Optional[int]							= [floats_list((1, x)							)[0] for x in range(800		,  1_400		,  200							)]
    _UpperCamelCase:     Optional[int]							= feature_extractor(
        _lowercase		,  padding='''longest'''		,  max_length=4		,  truncation=_lowercase		,  return_tensors='''np'''		,  return_attention_mask=_lowercase		,  )
    _UpperCamelCase:     Any							= inputs.input_features
    _UpperCamelCase:     Optional[int]							= inputs.attention_mask
    _UpperCamelCase:     int							= np.sum(attention_mask == 1		,  axis=1							)

    self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
    self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
    self._check_zero_mean_unit_variance(input_features[2]							)

    # make sure that if max_length < longest -> then pad to max_length
    self.assertEqual(input_features.shape		,  (3, 4, 24)							)

    _UpperCamelCase:     Tuple							= [floats_list((1, x)							)[0] for x in range(800		,  1_400		,  200							)]
    _UpperCamelCase:     Tuple							= feature_extractor(
        _lowercase		,  padding='''longest'''		,  max_length=16		,  truncation=_lowercase		,  return_tensors='''np'''		,  return_attention_mask=_lowercase		,  )
    _UpperCamelCase:     Union[str, Any]							= inputs.input_features
    _UpperCamelCase:     Optional[Any]							= inputs.attention_mask
    _UpperCamelCase:     Tuple							= np.sum(attention_mask == 1		,  axis=1							)

    self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
    self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
    self._check_zero_mean_unit_variance(input_features[2]							)

    # make sure that if max_length < longest -> then pad to max_length
    self.assertEqual(input_features.shape		,  (3, 6, 24)							)



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



    import torch

    _UpperCamelCase:     Tuple							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
    _UpperCamelCase:     int							= np.random.rand(100		,  32							).astype(np.floataa							)
    _UpperCamelCase:     Optional[int]							= np_speech_inputs.tolist()

    for inputs in [py_speech_inputs, np_speech_inputs]:
     _UpperCamelCase:     Optional[Any]							= feature_extractor.pad([{'''input_features''': inputs}]		,  return_tensors='''np'''							)
     self.assertTrue(np_processed.input_features.dtype == np.floataa							)
     _UpperCamelCase:     List[Any]							= feature_extractor.pad([{'''input_features''': inputs}]		,  return_tensors='''pt'''							)
     self.assertTrue(pt_processed.input_features.dtype == torch.floataa							)



   def   lowerCAmelCase						(	self					:  Any		,  _lowercase					:  Dict							):



    """simple docstring"""



    from datasets import load_dataset

    _UpperCamelCase:     Union[str, Any]							= load_dataset('''hf-internal-testing/librispeech_asr_dummy'''		,  '''clean'''		,  split='''validation'''							)
    # automatic decoding with librispeech
    _UpperCamelCase:     Optional[int]							= ds.sort('''id'''							).select(range(_lowercase							)							)[:num_samples]['''audio''']

    return [x["array"] for x in speech_samples]



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    _UpperCamelCase:     Union[str, Any]							= np.array([
        -1.5745, -1.7713, -1.7020, -1.6069, -1.2250, -1.1105, -0.9072, -0.8241,
        -1.2310, -0.8098, -0.3320, -0.4101, -0.7985, -0.4996, -0.8213, -0.9128,
        -1.0420, -1.1286, -1.0440, -0.7999, -0.8405, -1.2275, -1.5443, -1.4625,
    ]							)
    # fmt: on

    _UpperCamelCase:     Union[str, Any]							= self._load_datasamples(1							)
    _UpperCamelCase:     Union[str, Any]							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
    _UpperCamelCase:     int							= feature_extractor(_lowercase		,  return_tensors='''pt'''							).input_features
    self.assertEquals(input_features.shape		,  (1, 584, 24)							)
    self.assertTrue(np.allclose(input_features[0, 0, :30]		,  _lowercase		,  atol=1E-4							)							)
from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)




def       lowerCAmelCase_     (		lowercase:							Optional[Any]  ,					lowercase:							Tuple		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     Union[str, Any]							= []
 for part_id in partition_order:
  _UpperCamelCase:     int							= df.where(F"""SPARK_PARTITION_ID() = {part_id}"""		).collect()
  for row_idx, row in enumerate(lowercase		):
   expected_row_ids_and_row_dicts.append((F"""{part_id}_{row_idx}""", row.asDict())		)
 return expected_row_ids_and_row_dicts




@require_not_windows
@require_dill_gt_0_3_2
def       lowerCAmelCase_     (		)    ->				List[str]:







 '''simple docstring'''
 _UpperCamelCase:     Optional[int]							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     int							= spark.range(100		).repartition(1		)
 _UpperCamelCase:     int							= Spark(lowercase		)
 # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
 # that each partition can hold 2 rows.
 spark_builder._repartition_df_if_needed(max_shard_size=16		)
 # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
 assert spark_builder.df.rdd.getNumPartitions() == 50




@require_not_windows
@require_dill_gt_0_3_2
def       lowerCAmelCase_     (		)    ->				List[str]:







 '''simple docstring'''
 _UpperCamelCase:     List[Any]							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     Optional[Any]							= spark.range(10		).repartition(2		)
 _UpperCamelCase:     int							= [1, 0]
 _UpperCamelCase:     Any							= _generate_iterable_examples(lowercase  ,					lowercase		)  # Reverse the partitions.
 _UpperCamelCase:     List[str]							= _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase  ,					lowercase		)

 for i, (row_id, row_dict) in enumerate(generate_fn()		):
  _UpperCamelCase	,							_UpperCamelCase:     List[Any]							= expected_row_ids_and_row_dicts[i]
  assert row_id == expected_row_id
  assert row_dict == expected_row_dict




@require_not_windows
@require_dill_gt_0_3_2
def       lowerCAmelCase_     (		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     Optional[Any]							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     Any							= spark.range(10		).repartition(1		)
 _UpperCamelCase:     int							= SparkExamplesIterable(lowercase		)
 assert it.n_shards == 1
 for i, (row_id, row_dict) in enumerate(lowercase		):
  assert row_id == F"""0_{i}"""
  assert row_dict == {"id": i}




@require_not_windows
@require_dill_gt_0_3_2
def       lowerCAmelCase_     (		)    ->				Union[str, Any]:







 '''simple docstring'''
 _UpperCamelCase:     Any							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     str							= spark.range(30		).repartition(3		)
 # Mock the generator so that shuffle reverses the partition indices.
 with patch('''numpy.random.Generator'''		) as generator_mock:
  _UpperCamelCase:     Union[str, Any]							= lambda lowercase		: x.reverse()
  _UpperCamelCase:     List[str]							= _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase  ,					[2, 1, 0]		)

  _UpperCamelCase:     Union[str, Any]							= SparkExamplesIterable(lowercase		).shuffle_data_sources(lowercase		)
  assert shuffled_it.n_shards == 3
  for i, (row_id, row_dict) in enumerate(lowercase		):
   _UpperCamelCase	,							_UpperCamelCase:     List[Any]							= expected_row_ids_and_row_dicts[i]
   assert row_id == expected_row_id
   assert row_dict == expected_row_dict




@require_not_windows
@require_dill_gt_0_3_2
def       lowerCAmelCase_     (		)    ->				Union[str, Any]:







 '''simple docstring'''
 _UpperCamelCase:     str							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     Tuple							= spark.range(20		).repartition(4		)

 # Partitions 0 and 2
 _UpperCamelCase:     List[Any]							= SparkExamplesIterable(lowercase		).shard_data_sources(worker_id=0  ,					num_workers=2		)
 assert shard_it_a.n_shards == 2
 _UpperCamelCase:     Dict							= _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase  ,					[0, 2]		)
 for i, (row_id, row_dict) in enumerate(lowercase		):
  _UpperCamelCase	,							_UpperCamelCase:     Optional[int]							= expected_row_ids_and_row_dicts_a[i]
  assert row_id == expected_row_id
  assert row_dict == expected_row_dict

 # Partitions 1 and 3
 _UpperCamelCase:     str							= SparkExamplesIterable(lowercase		).shard_data_sources(worker_id=1  ,					num_workers=2		)
 assert shard_it_a.n_shards == 2
 _UpperCamelCase:     str							= _get_expected_row_ids_and_row_dicts_for_partition_order(lowercase  ,					[1, 3]		)
 for i, (row_id, row_dict) in enumerate(lowercase		):
  _UpperCamelCase	,							_UpperCamelCase:     str							= expected_row_ids_and_row_dicts_a[i]
  assert row_id == expected_row_id
  assert row_dict == expected_row_dict





@require_not_windows
@require_dill_gt_0_3_2
def       lowerCAmelCase_     (		)    ->				Optional[int]:







 '''simple docstring'''
 _UpperCamelCase:     Any							= pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
 _UpperCamelCase:     Union[str, Any]							= spark.range(100		).repartition(1		)
 _UpperCamelCase:     Optional[Any]							= Spark(lowercase		)
 # Choose a small max_shard_size for maximum partitioning.
 spark_builder._repartition_df_if_needed(max_shard_size=1		)
 # The new number of partitions should not be greater than the number of rows.
 assert spark_builder.df.rdd.getNumPartitions() == 100