
'''simple docstring'''







import numpy as np

from transformers import BatchFeature
from transformers.testing_utils import require_tf, require_torch

from .test_feature_extraction_common import FeatureExtractionSavingTestMixin


class 						__UpperCAmelCase      (					_lowerCamelCase						):
							# to overwrite at feature extractactor specific tests
							__lowercase      		=       None
							__lowercase      		=       None


							@property
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								return self.feat_extract_tester.prepare_feat_extract_dict()


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								self.assertTrue(hasattr(lowerCAmelCase_    ,   'feature_size'  )  )
								self.assertTrue(hasattr(lowerCAmelCase_    ,   'sampling_rate'  )  )
								self.assertTrue(hasattr(lowerCAmelCase_    ,   'padding_value'  )  )


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common()
								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}  )

								self.assertTrue(all(len(lowerCAmelCase_  ) == len(lowerCAmelCase_  ) for x, y in zip(lowerCAmelCase_    ,   processed_features[input_name]  )  )  )

								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common(equal_length=lowerCAmelCase_  )
								_snake_case   						=		BatchFeature({input_name: speech_inputs}    ,   tensor_type='np'  )

								_snake_case   						=		processed_features[input_name]

								if len(batch_features_input.shape  ) < 3:
									_snake_case   						=		batch_features_input[:, :, None]

								self.assertTrue(
								    batch_features_input.shape
								    == (self.feat_extract_tester.batch_size, len(speech_inputs[0]  ), self.feat_extract_tester.feature_size)  )


							@require_torch
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common(equal_length=lowerCAmelCase_  )
								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}    ,   tensor_type='pt'  )

								_snake_case   						=		processed_features[input_name]

								if len(batch_features_input.shape  ) < 3:
									_snake_case   						=		batch_features_input[:, :, None]

								self.assertTrue(
								    batch_features_input.shape
								    == (self.feat_extract_tester.batch_size, len(speech_inputs[0]  ), self.feat_extract_tester.feature_size)  )


							@require_tf
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common(equal_length=lowerCAmelCase_  )
								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}    ,   tensor_type='tf'  )

								_snake_case   						=		processed_features[input_name]

								if len(batch_features_input.shape  ) < 3:
									_snake_case   						=		batch_features_input[:, :, None]

								self.assertTrue(
								    batch_features_input.shape
								    == (self.feat_extract_tester.batch_size, len(speech_inputs[0]  ), self.feat_extract_tester.feature_size)  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_=False  ):





								"""simple docstring"""






								def _inputs_have_equal_length(lowerCAmelCase_  ):
									_snake_case   						=		len(input[0]  )
									for input_slice in input[1:]:
										if len(lowerCAmelCase_  ) != length:
											return False
									return True

								def _inputs_are_equal(lowerCAmelCase_    ,   lowerCAmelCase_  ):
									if len(lowerCAmelCase_  ) != len(lowerCAmelCase_  ):
										return False

									for input_slice_a, input_slice_a in zip(lowerCAmelCase_    ,   lowerCAmelCase_  ):
										if not np.allclose(np.asarray(lowerCAmelCase_  )    ,   np.asarray(lowerCAmelCase_  )    ,   atol=1E-3  ):
											return False
									return True

								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common(numpify=lowerCAmelCase_  )
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}  )

								_snake_case   						=		self.feat_extract_tester.seq_length_diff
								_snake_case   						=		self.feat_extract_tester.max_seq_length + pad_diff
								_snake_case   						=		self.feat_extract_tester.min_seq_length
								_snake_case   						=		self.feat_extract_tester.batch_size
								_snake_case   						=		self.feat_extract_tester.feature_size

								# test padding for List[int] + numpy
								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding=lowerCAmelCase_  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[-1]  )  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   return_tensors='np'  )
								_snake_case   						=		input_a[input_name]

								# max_length parameter has to be provided when setting `padding="max_length"`
								with self.assertRaises(lowerCAmelCase_  ):
									feat_extract.pad(lowerCAmelCase_    ,   padding='max_length'  )[input_name]

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=lowerCAmelCase_    ,   return_tensors='np'  )
								_snake_case   						=		input_a[input_name]

								self.assertFalse(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertTrue(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertTrue(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertTrue(_inputs_are_equal(lowerCAmelCase_    ,   lowerCAmelCase_  )  )
								self.assertTrue(len(input_a[0]  ) == pad_min_length  )
								self.assertTrue(len(input_a[1]  ) == pad_min_length + pad_diff  )
								self.assertTrue(input_a.shape[:2] == (batch_size, len(input_a[0]  ))  )
								self.assertTrue(input_a.shape[:2] == (batch_size, pad_max_length)  )

								if feature_size > 1:
									self.assertTrue(input_a.shape[2] == input_a.shape[2] == feature_size  )

								# test padding for `pad_to_multiple_of` for List[int] + numpy
								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   pad_to_multiple_of=10  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   pad_to_multiple_of=10  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   pad_to_multiple_of=10    ,   max_length=lowerCAmelCase_  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   pad_to_multiple_of=10    ,   max_length=lowerCAmelCase_    ,   return_tensors='np'    ,   )
								_snake_case   						=		input_a[input_name]

								self.assertTrue(all(len(lowerCAmelCase_  ) % 10 == 0 for x in input_a  )  )
								self.assertTrue(_inputs_are_equal(lowerCAmelCase_    ,   lowerCAmelCase_  )  )

								_snake_case   						=		pad_max_length if pad_max_length % 10 == 0 else (pad_max_length // 10 + 1) * 10
								self.assertTrue(all(len(lowerCAmelCase_  ) == expected_mult_pad_length for x in input_a  )  )
								self.assertEqual(input_a.shape[:2]    ,   (batch_size, expected_mult_pad_length)  )

								if feature_size > 1:
									self.assertTrue(input_a.shape[2] == feature_size  )

								# Check padding value is correct
								_snake_case   						=		(np.ones(self.feat_extract_tester.feature_size  ) * feat_extract.padding_value).sum()
								self.assertTrue(
								    abs(np.asarray(input_a[0]  )[pad_min_length:].sum() - padding_vector_sum * (pad_max_length - pad_min_length)  )
								    < 1E-3  )
								self.assertTrue(
								    abs(
								        np.asarray(input_a[1]  )[pad_min_length + pad_diff :].sum()
								        - padding_vector_sum * (pad_max_length - pad_min_length - pad_diff)  )
								    < 1E-3  )
								self.assertTrue(
								    abs(
								        np.asarray(input_a[2]  )[pad_min_length + 2 * pad_diff :].sum()
								        - padding_vector_sum * (pad_max_length - pad_min_length - 2 * pad_diff)  )
								    < 1E-3  )
								self.assertTrue(
								    abs(input_a[0, pad_min_length:].sum() - padding_vector_sum * (pad_max_length - pad_min_length)  ) < 1E-3  )
								self.assertTrue(
								    abs(input_a[0, pad_min_length:].sum() - padding_vector_sum * (expected_mult_pad_length - pad_min_length)  )
								    < 1E-3  )


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_=False  ):





								"""simple docstring"""






								def _inputs_have_equal_length(lowerCAmelCase_  ):
									_snake_case   						=		len(input[0]  )
									for input_slice in input[1:]:
										if len(lowerCAmelCase_  ) != length:
											return False
									return True

								def _inputs_are_equal(lowerCAmelCase_    ,   lowerCAmelCase_  ):
									if len(lowerCAmelCase_  ) != len(lowerCAmelCase_  ):
										return False

									for input_slice_a, input_slice_a in zip(lowerCAmelCase_    ,   lowerCAmelCase_  ):
										if not np.allclose(np.asarray(lowerCAmelCase_  )    ,   np.asarray(lowerCAmelCase_  )    ,   atol=1E-3  ):
											return False
									return True

								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common(numpify=lowerCAmelCase_  )
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}  )

								# truncate to smallest
								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[0]  )    ,   truncation=lowerCAmelCase_  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[0]  )  )
								_snake_case   						=		input_a[input_name]

								self.assertTrue(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertFalse(_inputs_have_equal_length(lowerCAmelCase_  )  )

								# truncate to smallest with np
								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[0]  )    ,   return_tensors='np'    ,   truncation=lowerCAmelCase_    ,   )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[0]  )    ,   return_tensors='np'  )
								_snake_case   						=		input_a[input_name]

								self.assertTrue(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertTrue(input_a.shape[1] == len(speech_inputs[0]  )  )

								# since truncation forces padding to be smaller than longest input
								# function can't return `np.ndarray`, but has to return list
								self.assertFalse(_inputs_have_equal_length(lowerCAmelCase_  )  )

								# truncate to middle
								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[1]  )    ,   truncation=lowerCAmelCase_    ,   return_tensors='np'    ,   )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[1]  )    ,   truncation=lowerCAmelCase_  )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[1]  )    ,   return_tensors='np'  )
								_snake_case   						=		input_a[input_name]

								self.assertTrue(input_a.shape[1] == len(speech_inputs[1]  )  )
								self.assertTrue(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertTrue(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertTrue(_inputs_are_equal(lowerCAmelCase_    ,   lowerCAmelCase_  )  )

								# since truncation forces padding to be smaller than longest input
								# function can't return `np.ndarray`, but has to return list
								self.assertFalse(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertTrue(len(input_a[-1]  ) == len(speech_inputs[-1]  )  )

								# padding has to be max_length when setting `truncation=True`
								with self.assertRaises(lowerCAmelCase_  ):
									feat_extract.pad(lowerCAmelCase_    ,   truncation=lowerCAmelCase_  )[input_name]

								# padding has to be max_length when setting `truncation=True`
								with self.assertRaises(lowerCAmelCase_  ):
									feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   truncation=lowerCAmelCase_  )[input_name]

								# padding has to be max_length when setting `truncation=True`
								with self.assertRaises(lowerCAmelCase_  ):
									feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   truncation=lowerCAmelCase_  )[input_name]

								# max_length parameter has to be provided when setting `truncation=True` and padding="max_length"
								with self.assertRaises(lowerCAmelCase_  ):
									feat_extract.pad(lowerCAmelCase_    ,   padding='max_length'    ,   truncation=lowerCAmelCase_  )[input_name]

								# test truncation for `pad_to_multiple_of` for List[int] + numpy
								_snake_case   						=		12
								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[0]  )    ,   pad_to_multiple_of=lowerCAmelCase_    ,   truncation=lowerCAmelCase_    ,   )
								_snake_case   						=		input_a[input_name]

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=len(speech_inputs[0]  )    ,   pad_to_multiple_of=lowerCAmelCase_    ,   )
								_snake_case   						=		input_a[input_name]

								# retrieve expected_length as multiple of pad_to_multiple_of
								_snake_case   						=		len(speech_inputs[0]  )
								if expected_length % pad_to_multiple_of != 0:
									_snake_case   						=		((len(speech_inputs[0]  ) // pad_to_multiple_of) + 1) * pad_to_multiple_of

								self.assertTrue(len(input_a[0]  ) == expected_length  )
								self.assertTrue(_inputs_have_equal_length(lowerCAmelCase_  )  )
								self.assertFalse(_inputs_have_equal_length(lowerCAmelCase_  )  )


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								self._check_padding(numpify=lowerCAmelCase_  )


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								self._check_padding(numpify=lowerCAmelCase_  )


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								self._check_truncation(numpify=lowerCAmelCase_  )


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								self._check_truncation(numpify=lowerCAmelCase_  )


							@require_torch
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common()
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}  )

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   return_tensors='np'  )[input_name]
								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   return_tensors='pt'  )[input_name]

								self.assertTrue(abs(input_np.astype(np.floataa  ).sum() - input_pt.numpy().astype(np.floataa  ).sum()  ) < 1E-2  )


							@require_tf
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feature_extraction_class(**self.feat_extract_dict  )
								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common()
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}  )

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   return_tensors='np'  )[input_name]
								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   return_tensors='tf'  )[input_name]

								self.assertTrue(abs(input_np.astype(np.floataa  ).sum() - input_tf.numpy().astype(np.floataa  ).sum()  ) < 1E-2  )


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feat_extract_dict
								_snake_case   						=		True
								_snake_case   						=		self.feature_extraction_class(**lowerCAmelCase_  )
								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common()
								_snake_case   						=		[len(lowerCAmelCase_  ) for x in speech_inputs]
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}  )

								_snake_case   						=		feat_extract.pad(lowerCAmelCase_    ,   padding='longest'    ,   return_tensors='np'  )
								self.assertIn('attention_mask'    ,   lowerCAmelCase_  )
								self.assertListEqual(list(processed.attention_mask.shape  )    ,   list(processed[input_name].shape[:2]  )  )
								self.assertListEqual(processed.attention_mask.sum(-1  ).tolist()    ,   lowerCAmelCase_  )




							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		self.feat_extract_dict
								_snake_case   						=		True
								_snake_case   						=		self.feature_extraction_class(**lowerCAmelCase_  )
								_snake_case   						=		self.feat_extract_tester.prepare_inputs_for_common()
								_snake_case   						=		[len(lowerCAmelCase_  ) for x in speech_inputs]
								_snake_case   						=		feat_extract.model_input_names[0]

								_snake_case   						=		BatchFeature({input_name: speech_inputs}  )
								_snake_case   						=		min(lowerCAmelCase_  )

								_snake_case   						=		feat_extract.pad(
								    lowerCAmelCase_    ,   padding='max_length'    ,   max_length=lowerCAmelCase_    ,   truncation=lowerCAmelCase_    ,   return_tensors='np'  )
								self.assertIn('attention_mask'    ,   lowerCAmelCase_  )
								self.assertListEqual(
								    list(processed_pad.attention_mask.shape  )    ,   [processed_pad[input_name].shape[0], max_length]  )
								self.assertListEqual(
								    processed_pad.attention_mask[:, :max_length].sum(-1  ).tolist()    ,   [max_length for x in speech_inputs]  )









from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				def __init__(     self						:    Dict,  lowerCamelCase						:    pyspark.sql.DataFrame,  lowerCamelCase						:    Optional[NamedSplit] = None,  lowerCamelCase						:    Optional[Features] = None,  lowerCamelCase						:    bool = True,  lowerCamelCase						:    str = None,  lowerCamelCase						:    bool = False,  lowerCamelCase						:    str = None,  lowerCamelCase						:    bool = True,  lowerCamelCase						:    str = "arrow",  **lowerCamelCase						:    str,  ):

								'''simple docstring'''
								super().__init__(
								    split=lowerCamelCase,  features=lowerCamelCase,  cache_dir=lowerCamelCase,  keep_in_memory=lowerCamelCase,  streaming=lowerCamelCase,  **lowerCamelCase,  )
								lowercase__    =			load_from_cache_file
								lowercase__    =			file_format
								lowercase__    =			Spark(
								    df=lowerCamelCase,  features=lowerCamelCase,  cache_dir=lowerCamelCase,  working_dir=lowerCamelCase,  **lowerCamelCase,  )




				def 				lowercase__   (     self						:    List[Any]  ):

								'''simple docstring'''
								if self.streaming:
												return self.builder.as_streaming_dataset(split=self.split  )
								lowercase__    =			None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
								self.builder.download_and_prepare(
								    download_mode=lowerCamelCase,  file_format=self._file_format,  )
								return self.builder.as_dataset(split=self.split  )




