




import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
    import jax.numpy as jnp

if is_tf_available():
    import tensorflow as tf

if is_torch_available():
    import torch
class  __a   (					unittest.TestCase	):


     """simple docstring"""


     def     __A							(      self     :					int		)						->  Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        ={
             """task_specific_params""": {
                 """summarization""": {"""length_penalty""": 1.0, """max_length""": 1_2_8, """min_length""": 1_2, """num_beams""": 4},
                 """summarization_cnn""": {"""length_penalty""": 2.0, """max_length""": 1_4_2, """min_length""": 5_6, """num_beams""": 4},
                 """summarization_xsum""": {"""length_penalty""": 1.0, """max_length""": 6_2, """min_length""": 1_1, """num_beams""": 6},
             }
         }
         SCREAMING_SNAKE_CASE__        ={
             """task_specific_params.summarization.length_penalty""": 1.0,
             """task_specific_params.summarization.max_length""": 1_2_8,
             """task_specific_params.summarization.min_length""": 1_2,
             """task_specific_params.summarization.num_beams""": 4,
             """task_specific_params.summarization_cnn.length_penalty""": 2.0,
             """task_specific_params.summarization_cnn.max_length""": 1_4_2,
             """task_specific_params.summarization_cnn.min_length""": 5_6,
             """task_specific_params.summarization_cnn.num_beams""": 4,
             """task_specific_params.summarization_xsum.length_penalty""": 1.0,
             """task_specific_params.summarization_xsum.max_length""": 6_2,
             """task_specific_params.summarization_xsum.min_length""": 1_1,
             """task_specific_params.summarization_xsum.num_beams""": 6,
         }

         self.assertEqual(flatten_dict(__lowerCAmelCase		)							,__lowerCAmelCase		)




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


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase		)							,x.transpose()		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,axes=(1, 2, 0)		)							,x.transpose((1, 2, 0)		)		)		)




     @require_torch
     def     __A							(      self     :					int		)						->  List[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =torch.tensor(__lowerCAmelCase		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase		)							,transpose(__lowerCAmelCase		).numpy()		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         SCREAMING_SNAKE_CASE__        =torch.tensor(__lowerCAmelCase		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,axes=(1, 2, 0)		)							,transpose(__lowerCAmelCase							,axes=(1, 2, 0)		).numpy()		)		)




     @require_tf
     def     __A							(      self     :					Dict		)						->  List[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =tf.constant(__lowerCAmelCase		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase		)							,transpose(__lowerCAmelCase		).numpy()		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         SCREAMING_SNAKE_CASE__        =tf.constant(__lowerCAmelCase		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,axes=(1, 2, 0)		)							,transpose(__lowerCAmelCase							,axes=(1, 2, 0)		).numpy()		)		)




     @require_flax
     def     __A							(      self     :					str		)						->  List[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =jnp.array(__lowerCAmelCase		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase		)							,np.asarray(transpose(__lowerCAmelCase		)		)		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         SCREAMING_SNAKE_CASE__        =jnp.array(__lowerCAmelCase		)
         self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,axes=(1, 2, 0)		)							,np.asarray(transpose(__lowerCAmelCase							,axes=(1, 2, 0)		)		)		)		)




     def     __A							(      self     :					Tuple		)						->  Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(4, 3)		)							,np.reshape(__lowerCAmelCase							,(4, 3)		)		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(1_2, 5)		)							,np.reshape(__lowerCAmelCase							,(1_2, 5)		)		)		)




     @require_torch
     def     __A							(      self     :					Tuple		)						->  Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =torch.tensor(__lowerCAmelCase		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(4, 3)		)							,reshape(__lowerCAmelCase							,(4, 3)		).numpy()		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         SCREAMING_SNAKE_CASE__        =torch.tensor(__lowerCAmelCase		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(1_2, 5)		)							,reshape(__lowerCAmelCase							,(1_2, 5)		).numpy()		)		)




     @require_tf
     def     __A							(      self     :					List[Any]		)						->  Dict:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =tf.constant(__lowerCAmelCase		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(4, 3)		)							,reshape(__lowerCAmelCase							,(4, 3)		).numpy()		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         SCREAMING_SNAKE_CASE__        =tf.constant(__lowerCAmelCase		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(1_2, 5)		)							,reshape(__lowerCAmelCase							,(1_2, 5)		).numpy()		)		)




     @require_flax
     def     __A							(      self     :					List[str]		)						->  Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =jnp.array(__lowerCAmelCase		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(4, 3)		)							,np.asarray(reshape(__lowerCAmelCase							,(4, 3)		)		)		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4							,5		)
         SCREAMING_SNAKE_CASE__        =jnp.array(__lowerCAmelCase		)
         self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,(1_2, 5)		)							,np.asarray(reshape(__lowerCAmelCase							,(1_2, 5)		)		)		)		)




     def     __A							(      self     :					List[str]		)						->  Optional[int]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,3							,4		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase		)							,np.squeeze(__lowerCAmelCase		)		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,4							,1							,5		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,axis=2		)							,np.squeeze(__lowerCAmelCase							,axis=2		)		)		)




     @require_torch
     def     __A							(      self     :					Any		)						->  int:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,3							,4		)
         SCREAMING_SNAKE_CASE__        =torch.tensor(__lowerCAmelCase		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase		)							,squeeze(__lowerCAmelCase		).numpy()		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,4							,1							,5		)
         SCREAMING_SNAKE_CASE__        =torch.tensor(__lowerCAmelCase		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,axis=2		)							,squeeze(__lowerCAmelCase							,axis=2		).numpy()		)		)




     @require_tf
     def     __A							(      self     :					int		)						->  Optional[int]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,3							,4		)
         SCREAMING_SNAKE_CASE__        =tf.constant(__lowerCAmelCase		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase		)							,squeeze(__lowerCAmelCase		).numpy()		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,4							,1							,5		)
         SCREAMING_SNAKE_CASE__        =tf.constant(__lowerCAmelCase		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,axis=2		)							,squeeze(__lowerCAmelCase							,axis=2		).numpy()		)		)




     @require_flax
     def     __A							(      self     :					int		)						->  Union[str, Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,3							,4		)
         SCREAMING_SNAKE_CASE__        =jnp.array(__lowerCAmelCase		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase		)							,np.asarray(squeeze(__lowerCAmelCase		)		)		)		)

         SCREAMING_SNAKE_CASE__        =np.random.randn(1							,4							,1							,5		)
         SCREAMING_SNAKE_CASE__        =jnp.array(__lowerCAmelCase		)
         self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,axis=2		)							,np.asarray(squeeze(__lowerCAmelCase							,axis=2		)		)		)		)




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


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,axis=1		)							,np.expand_dims(__lowerCAmelCase							,axis=1		)		)		)




     @require_torch
     def     __A							(      self     :					List[Any]		)						->  Any:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =torch.tensor(__lowerCAmelCase		)
         self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,axis=1		)							,expand_dims(__lowerCAmelCase							,axis=1		).numpy()		)		)




     @require_tf
     def     __A							(      self     :					int		)						->  Optional[Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =tf.constant(__lowerCAmelCase		)
         self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,axis=1		)							,expand_dims(__lowerCAmelCase							,axis=1		).numpy()		)		)




     @require_flax
     def     __A							(      self     :					Optional[Any]		)						->  Union[str, Any]:


         '''simple docstring'''





         SCREAMING_SNAKE_CASE__        =np.random.randn(3							,4		)
         SCREAMING_SNAKE_CASE__        =jnp.array(__lowerCAmelCase		)
         self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,axis=1		)							,np.asarray(expand_dims(__lowerCAmelCase							,axis=1		)		)		)		)













import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser







def   __lowerCamelCase   (				__a  :Optional[int]       )						->							str:






		"""simple docstring"""

		A__     =  {}
		A__     =  tokenizer(example["""content"""]       ,	truncation=__a       )["""input_ids"""]
		A__     =  len(example["""content"""]       ) / len(output["""input_ids"""]       )
		return output


A       :   Union[str, Any]					= HfArgumentParser(PretokenizationArguments)
A       :   str					= parser.parse_args()
if args.num_workers is None:
					A       :   Any					= multiprocessing.cpu_count()
A       :   int					= AutoTokenizer.from_pretrained(args.tokenizer_dir)

A       :   Optional[int]					= time.time()
A       :   int					= load_dataset(args.dataset_name, split='''train''')
print(F'''Dataset loaded in {time.time()-t_start:.2f}s''')

A       :   Optional[Any]					= time.time()
A       :   int					= 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       :   Tuple					= time.time()
ds.push_to_hub(args.tokenized_data_repo)






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



