


"""simple docstring"""







import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam






class 				__A    (	datasets.BeamBasedBuilder      ):






 """simple docstring"""




 def 				SCREAMING_SNAKE_CASE							(		self						) ->					str:
     return datasets.DatasetInfo(
         features=datasets.Features({'''content''': datasets.Value('''string'''						)}						)			,   supervised_keys=__A			,   )


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A						) ->					List[str]:
     return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,   gen_kwargs={'''examples''': get_test_dummy_examples()}						)]





 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A						) ->					Optional[int]:
     import apache_beam as beam

     return pipeline | "Load Examples" >> beam.Create(__A						)






class 				__A    (	datasets.BeamBasedBuilder      ):






 """simple docstring"""




 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
     return datasets.DatasetInfo(
         features=datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string'''						)}						)}						)			,   supervised_keys=__A			,   )


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A						) ->					Optional[Any]:
     return [
         datasets.SplitGenerator(name=datasets.Split.TRAIN			,   gen_kwargs={'''examples''': get_test_nested_examples()}						)
     ]





 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A						) ->					int:
     import apache_beam as beam

     return pipeline | "Load Examples" >> beam.Create(__A						)

def _A						(						):







    """simple docstring"""
    return [(i, {"content": content}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']		)]
def _A						(						):







    """simple docstring"""
    return [(i, {"a": {"b": [content]}}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']		)]






class 				__A    (	SCREAMING_SNAKE_CASE__      ):






 """simple docstring"""




 @require_beam
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					int:
     a					=len(get_test_dummy_examples()						)
     with tempfile.TemporaryDirectory() as tmp_cache_dir:
         a					=DummyBeamDataset(cache_dir=__A			,   beam_runner='''DirectRunner'''						)
         builder.download_and_prepare()
         self.assertTrue(
             os.path.exists(
                 os.path.join(__A			,   builder.name			,   '''default'''			,   '''0.0.0'''			,   f'''{builder.name}-train.arrow'''						)						)						)
         self.assertDictEqual(builder.info.features			,   datasets.Features({'''content''': datasets.Value('''string'''						)}						)						)
         a					=builder.as_dataset()
         self.assertEqual(dset['''train'''].num_rows			,   __A						)
         self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples			,   __A						)
         self.assertDictEqual(dset['''train'''][0]			,   get_test_dummy_examples()[0][1]						)
         self.assertDictEqual(
             dset['''train'''][expected_num_examples - 1]			,   get_test_dummy_examples()[expected_num_examples - 1][1]						)
         self.assertTrue(
             os.path.exists(os.path.join(__A			,   builder.name			,   '''default'''			,   '''0.0.0'''			,   '''dataset_info.json'''						)						)						)
         del dset


 @require_beam
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     import apache_beam as beam

     a					=beam.io.parquetio.WriteToParquet

     a					=len(get_test_dummy_examples()						)
     with tempfile.TemporaryDirectory() as tmp_cache_dir:
         a					=DummyBeamDataset(cache_dir=__A			,   beam_runner='''DirectRunner'''						)
         with patch('''apache_beam.io.parquetio.WriteToParquet'''						) as write_parquet_mock:
             a					=partial(__A			,   num_shards=2						)
             builder.download_and_prepare()
         self.assertTrue(
             os.path.exists(
                 os.path.join(
                     __A			,   builder.name			,   '''default'''			,   '''0.0.0'''			,   f'''{builder.name}-train-00000-of-00002.arrow'''						)						)						)
         self.assertTrue(
             os.path.exists(
                 os.path.join(
                     __A			,   builder.name			,   '''default'''			,   '''0.0.0'''			,   f'''{builder.name}-train-00000-of-00002.arrow'''						)						)						)
         self.assertDictEqual(builder.info.features			,   datasets.Features({'''content''': datasets.Value('''string'''						)}						)						)
         a					=builder.as_dataset()
         self.assertEqual(dset['''train'''].num_rows			,   __A						)
         self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples			,   __A						)
         # Order is not preserved when sharding, so we just check that all the elements are there
         self.assertListEqual(sorted(dset['''train''']['''content''']						)			,   sorted(['''foo''', '''bar''', '''foobar''']						)						)
         self.assertTrue(
             os.path.exists(os.path.join(__A			,   builder.name			,   '''default'''			,   '''0.0.0'''			,   '''dataset_info.json'''						)						)						)
         del dset


 @require_beam
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					int:
     with tempfile.TemporaryDirectory() as tmp_cache_dir:
         a					=DummyBeamDataset(cache_dir=__A						)
         self.assertRaises(datasets.builder.MissingBeamOptions			,   builder.download_and_prepare						)





 @require_beam
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Dict:
     a					=len(get_test_nested_examples()						)
     with tempfile.TemporaryDirectory() as tmp_cache_dir:
         a					=NestedBeamDataset(cache_dir=__A			,   beam_runner='''DirectRunner'''						)
         builder.download_and_prepare()
         self.assertTrue(
             os.path.exists(
                 os.path.join(__A			,   builder.name			,   '''default'''			,   '''0.0.0'''			,   f'''{builder.name}-train.arrow'''						)						)						)
         self.assertDictEqual(
             builder.info.features			,   datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string'''						)}						)}						)						)
         a					=builder.as_dataset()
         self.assertEqual(dset['''train'''].num_rows			,   __A						)
         self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples			,   __A						)
         self.assertDictEqual(dset['''train'''][0]			,   get_test_nested_examples()[0][1]						)
         self.assertDictEqual(
             dset['''train'''][expected_num_examples - 1]			,   get_test_nested_examples()[expected_num_examples - 1][1]						)
         self.assertTrue(
             os.path.exists(os.path.join(__A			,   builder.name			,   '''default'''			,   '''0.0.0'''			,   '''dataset_info.json'''						)						)						)
         del dset



"""simple docstring"""







import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser
def _A						(						lowercase		):







    """simple docstring"""
    a					={}
    a					=tokenizer(example['''content''']	,   truncation=lowercase		)['''input_ids''']
    a					=len(example['''content''']		) / len(output['''input_ids''']		)
    return output


lowerCamelCase_       :       Optional[int]      						=  HfArgumentParser(PretokenizationArguments)
lowerCamelCase_       :       Optional[Any]      						=  parser.parse_args()
if args.num_workers is None:
       lowerCamelCase_       :       Tuple      						=  multiprocessing.cpu_count()
lowerCamelCase_       :       Any      						=  AutoTokenizer.from_pretrained(args.tokenizer_dir)

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

lowerCamelCase_       :       List[str]      						=  time.time()
lowerCamelCase_       :       str      						=  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')

lowerCamelCase_       :       Union[str, Any]      						=  time.time()
ds.push_to_hub(args.tokenized_data_repo)





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