


'''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   __lowercase				(  __a  ):
      UpperCamelCase								=       None
      UpperCamelCase								=       None





      @property
      def        _lowercase				(		self						:					List[str]       )    ->   Union[str, Any]:






         """simple docstring"""







         return self.feat_extract_tester.prepare_feat_extract_dict()





      def        _lowercase				(		self						:					Optional[Any]       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     self.feature_extraction_class(**self.feat_extract_dict       )
         self.assertTrue(hasattr(a_					,       """feature_size"""       )       )
         self.assertTrue(hasattr(a_					,       """sampling_rate"""       )       )
         self.assertTrue(hasattr(a_					,       """padding_value"""       )       )





      def        _lowercase				(		self						:					str       )    ->   Dict:






         """simple docstring"""







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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}       )

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

         UpperCAmelCase							=     self.feat_extract_tester.prepare_inputs_for_common(equal_length=a_       )
         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}					,       tensor_type="""np"""       )

         UpperCAmelCase							=     processed_features[input_name]

         if len(batch_features_input.shape       ) < 3:
            UpperCAmelCase							=     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        _lowercase				(		self						:					int       )    ->   Dict:






         """simple docstring"""







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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}					,       tensor_type="""pt"""       )

         UpperCAmelCase							=     processed_features[input_name]

         if len(batch_features_input.shape       ) < 3:
            UpperCAmelCase							=     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        _lowercase				(		self						:					Union[str, Any]       )    ->   str:






         """simple docstring"""







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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}					,       tensor_type="""tf"""       )

         UpperCAmelCase							=     processed_features[input_name]

         if len(batch_features_input.shape       ) < 3:
            UpperCAmelCase							=     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        _lowercase				(		self						:					List[str]					,       __lowerCamelCase						:					Union[str, Any]=False       )    ->   Dict:






         """simple docstring"""







         def _inputs_have_equal_length(__lowerCamelCase						:					Union[str, Any]       ):
            UpperCAmelCase							=     len(input[0]       )
            for input_slice in input[1:]:
               if len(a_       ) != length:
                  return False
            return True

         def _inputs_are_equal(__lowerCamelCase						:					List[str]					,       __lowerCamelCase						:					int       ):
            if len(a_       ) != len(a_       ):
               return False

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

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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}       )

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

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

         UpperCAmelCase							=     feat_extract.pad(a_					,       padding="""longest"""       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[-1]       )       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(a_					,       padding="""longest"""					,       return_tensors="""np"""       )
         UpperCAmelCase							=     input_a[input_name]

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

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=a_					,       return_tensors="""np"""       )
         UpperCAmelCase							=     input_a[input_name]

         self.assertFalse(_inputs_have_equal_length(a_       )       )
         self.assertTrue(_inputs_have_equal_length(a_       )       )
         self.assertTrue(_inputs_have_equal_length(a_       )       )
         self.assertTrue(_inputs_are_equal(a_					,       a_       )       )
         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
         UpperCAmelCase							=     feat_extract.pad(a_					,       pad_to_multiple_of=1_0       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(a_					,       padding="""longest"""					,       pad_to_multiple_of=1_0       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       pad_to_multiple_of=1_0					,       max_length=a_       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       pad_to_multiple_of=1_0					,       max_length=a_					,       return_tensors="""np"""					,       )
         UpperCAmelCase							=     input_a[input_name]

         self.assertTrue(all(len(a_       ) % 1_0 == 0 for x in input_a       )       )
         self.assertTrue(_inputs_are_equal(a_					,       a_       )       )

         UpperCAmelCase							=     pad_max_length if pad_max_length % 1_0 == 0 else (pad_max_length // 1_0 + 1) * 1_0
         self.assertTrue(all(len(a_       ) == 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
         UpperCAmelCase							=     (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        _lowercase				(		self						:					Optional[Any]					,       __lowerCamelCase						:					Union[str, Any]=False       )    ->   str:






         """simple docstring"""







         def _inputs_have_equal_length(__lowerCamelCase						:					Union[str, Any]       ):
            UpperCAmelCase							=     len(input[0]       )
            for input_slice in input[1:]:
               if len(a_       ) != length:
                  return False
            return True

         def _inputs_are_equal(__lowerCamelCase						:					Tuple					,       __lowerCamelCase						:					Any       ):
            if len(a_       ) != len(a_       ):
               return False

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

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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}       )

         # truncate to smallest
         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[0]       )					,       truncation=a_       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[0]       )       )
         UpperCAmelCase							=     input_a[input_name]

         self.assertTrue(_inputs_have_equal_length(a_       )       )
         self.assertFalse(_inputs_have_equal_length(a_       )       )

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

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[0]       )					,       return_tensors="""np"""       )
         UpperCAmelCase							=     input_a[input_name]

         self.assertTrue(_inputs_have_equal_length(a_       )       )
         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(a_       )       )

         # truncate to middle
         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[1]       )					,       truncation=a_					,       return_tensors="""np"""					,       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[1]       )					,       truncation=a_       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[1]       )					,       return_tensors="""np"""       )
         UpperCAmelCase							=     input_a[input_name]

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

         # 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(a_       )       )
         self.assertTrue(len(input_a[-1]       ) == len(speech_inputs[-1]       )       )

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

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

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

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

         # test truncation for `pad_to_multiple_of` for List[int] + numpy
         UpperCAmelCase							=     1_2
         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[0]       )					,       pad_to_multiple_of=a_					,       truncation=a_					,       )
         UpperCAmelCase							=     input_a[input_name]

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=len(speech_inputs[0]       )					,       pad_to_multiple_of=a_					,       )
         UpperCAmelCase							=     input_a[input_name]

         # retrieve expected_length as multiple of pad_to_multiple_of
         UpperCAmelCase							=     len(speech_inputs[0]       )
         if expected_length % pad_to_multiple_of != 0:
            UpperCAmelCase							=     ((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(a_       )       )
         self.assertFalse(_inputs_have_equal_length(a_       )       )





      def        _lowercase				(		self						:					int       )    ->   List[Any]:






         """simple docstring"""







         self._check_padding(numpify=a_       )





      def        _lowercase				(		self						:					Dict       )    ->   Optional[int]:






         """simple docstring"""







         self._check_padding(numpify=a_       )





      def        _lowercase				(		self						:					Optional[int]       )    ->   List[Any]:






         """simple docstring"""







         self._check_truncation(numpify=a_       )





      def        _lowercase				(		self						:					Tuple       )    ->   Tuple:






         """simple docstring"""







         self._check_truncation(numpify=a_       )





      @require_torch
      def        _lowercase				(		self						:					Union[str, Any]       )    ->   Optional[Any]:






         """simple docstring"""







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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}       )

         UpperCAmelCase							=     feat_extract.pad(a_					,       padding="""longest"""					,       return_tensors="""np"""       )[input_name]
         UpperCAmelCase							=     feat_extract.pad(a_					,       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        _lowercase				(		self						:					int       )    ->   Optional[int]:






         """simple docstring"""







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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}       )

         UpperCAmelCase							=     feat_extract.pad(a_					,       padding="""longest"""					,       return_tensors="""np"""       )[input_name]
         UpperCAmelCase							=     feat_extract.pad(a_					,       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        _lowercase				(		self						:					List[Any]       )    ->   Optional[int]:






         """simple docstring"""







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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}       )

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





      def        _lowercase				(		self						:					Optional[Any]       )    ->   Optional[int]:






         """simple docstring"""







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

         UpperCAmelCase							=     BatchFeature({input_name: speech_inputs}       )
         UpperCAmelCase							=     min(a_       )

         UpperCAmelCase							=     feat_extract.pad(
             a_					,       padding="""max_length"""					,       max_length=a_					,       truncation=a_					,       return_tensors="""np"""       )
         self.assertIn("""attention_mask"""					,       a_       )
         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]       )







import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser
def   _UpperCamelCase				(						lowerCAmelCase_   )						->int:
   UpperCAmelCase							=     {}
   UpperCAmelCase							=     tokenizer(example["""content"""]  ,						truncation=lowerCAmelCase_   )["""input_ids"""]
   UpperCAmelCase							=     len(example["""content"""]   ) / len(output["""input_ids"""]   )
   return output


__a     =			HfArgumentParser(PretokenizationArguments)
__a     =			parser.parse_args()
if args.num_workers is None:
 __a     =			multiprocessing.cpu_count()
__a     =			AutoTokenizer.from_pretrained(args.tokenizer_dir)

__a     =			time.time()
__a     =			load_dataset(args.dataset_name, split="""train""")
print(F"""Dataset loaded in {time.time()-t_start:.2f}s""")

__a     =			time.time()
__a     =			ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        """repo_name""",
        """path""",
        """copies""",
        """size""",
        """content""",
        """license""",
        """hash""",
        """line_mean""",
        """line_max""",
        """alpha_frac""",
        """autogenerated""",
    ],
)
print(F"""Dataset tokenized in {time.time()-t_start:.2f}s""")

__a     =			time.time()
ds.push_to_hub(args.tokenized_data_repo)







print(F"""Data pushed to the hub in {time.time()-t_start:.2f}s""")



