

'''simple docstring'''






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


__SCREAMING_SNAKE_CASE     =						logging.getLogger(__name__)



def 			__a		(       ):
     a__   :      Union[str, Any]     =						argparse.ArgumentParser(
         description='''Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.'''      )
     parser.add_argument(
         '''--dataset_name'''     ,  type=lowerCamelCase__     ,  default='''wikitext'''     ,  help='''Name of the training. Explore datasets at: hf.co/datasets.'''     ,  )
     parser.add_argument(
         '''--dataset_config'''     ,  type=lowerCamelCase__     ,  default='''wikitext-103-raw-v1'''     ,  help='''Configuration name of the dataset.'''      )
     parser.add_argument(
         '''--tokenizer_name_or_path'''     ,  type=lowerCamelCase__     ,  default='''sayakpaul/unigram-tokenizer-wikitext'''     ,  help='''Tokenizer identifier. Can be a local filepath or a Hub identifier.'''     ,  )
     parser.add_argument(
         '''--shard_size'''     ,  type=lowerCamelCase__     ,  default=1000     ,  help='''Number of entries to go in a single shard.'''     ,  )
     parser.add_argument('''--split'''     ,  type=lowerCamelCase__     ,  default='''train'''     ,  choices=['''train''', '''test''', '''validation''']      )
     parser.add_argument(
         '''--limit'''     ,  default=lowerCamelCase__     ,  type=lowerCamelCase__     ,  help='''Limit the number of shards (used for debugging).'''     ,  )
     parser.add_argument(
         '''--max_length'''     ,  type=lowerCamelCase__     ,  default=512     ,  help='''Maximum sequence length. For training on TPUs, it helps to have a maximum'''
         ''' sequence length that is a multiple of 8.'''     ,  )
     parser.add_argument(
         '''--output_dir'''     ,  default='''tf-tpu'''     ,  type=lowerCamelCase__     ,  help='''Output directory where the TFRecord shards will be saved. If the'''
         ''' path is appended with `gs://` (\'gs://tf-tpu\', for example) then the TFRecord'''
         ''' shards will be directly saved to a Google Cloud Storage bucket.'''     ,  )

     a__   :      Dict     =						parser.parse_args()
     return args



def 			__a		(       lowerCAmelCase__ :    Tuple      ):
     def fn(lowerCAmelCase__ :    List[Any]      ):
          return tokenizer(examples['''text''']      )

     return fn



def 			__a		(       lowerCAmelCase__ :    List[str]      ):
     a__   :      List[Any]     =						[]
     for i in range(len(tokenized_data['''input_ids''']      )      ):
          a__   :      Dict     =						{
              '''input_ids''': tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data['''input_ids'''][i]      )      ),
              '''attention_mask''': tf.train.Feature(
                  intaa_list=tf.train.IntaaList(value=tokenized_data['''attention_mask'''][i]      )      ),
          }
          a__   :      Optional[int]     =						tf.train.Features(feature=lowerCamelCase__      )
          a__   :      Dict     =						tf.train.Example(features=lowerCamelCase__      )
          a__   :      List[str]     =						example.SerializeToString()
          records.append(lowerCamelCase__      )
     return records



def 			__a		(       lowerCAmelCase__ :    List[Any]      ):
     a__   :      Tuple     =						datasets.load_dataset(args.dataset_name     ,  args.dataset_config     ,  split=args.split      )

     if args.limit is not None:
          a__   :      Union[str, Any]     =						min(len(lowerCamelCase__      )     ,  args.limit      )
          a__   :      int     =						dataset.select(range(lowerCamelCase__      )      )
          print(F'Limiting the dataset to {args.limit} entries.'      )

     a__   :      Tuple     =						AutoTokenizer.from_pretrained(args.tokenizer_name_or_path      )

     # Handle output directory creation.
     # For serializing into a Google Cloud Storage Bucket, one needs to first
     # create a bucket.
     if "gs" not in args.output_dir:
          if not os.path.exists(args.output_dir      ):
               os.makedirs(args.output_dir      )
          a__   :      Any     =						os.path.join(args.output_dir     ,  args.split      )
          if not os.path.exists(lowerCamelCase__      ):
               os.makedirs(lowerCamelCase__      )
     else:
          a__   :      int     =						os.path.join(args.output_dir     ,  args.split      )

     # Tokenize the whole dataset at once.
     a__   :      List[Any]     =						tokenize_function(lowerCamelCase__      )
     a__   :      str     =						dataset.map(lowerCamelCase__     ,  batched=lowerCamelCase__     ,  num_proc=4     ,  remove_columns=['''text''']      )

     # We need to concatenate all our texts together, and then split the result
     # into chunks of a fixed size, which we will call block_size. To do this, we
     # will use the map method again, with the option batched=True. When we use batched=True,
     # the function we pass to map() will be passed multiple inputs at once, allowing us
     # to group them into more or fewer examples than we had in the input.
     # This allows us to create our new fixed-length samples. The advantage of this
     # method is that we don't lose a whole lot of content from the dataset compared to the
     # case where we simply tokenize with a pre-defined max_length.

     def group_texts(lowerCAmelCase__ :    Dict      ):
          # Concatenate all texts.
          a__   :      Optional[int]     =						{k: sum(examples[k]     ,  []      ) for k in examples.keys()}
          a__   :      Optional[int]     =						len(concatenated_examples[list(examples.keys()      )[0]]      )
          # We drop the small remainder, though you could add padding instead if the model supports it
          # In this, as in all things, we advise you to follow your heart 🫀
          a__   :      List[Any]     =						(total_length // args.max_length) * args.max_length
          # Split by chunks of max_len.
          a__   :      Optional[int]     =						{
              k: [t[i : i + args.max_length] for i in range(0     ,  lowerCamelCase__     ,  args.max_length      )]
              for k, t in concatenated_examples.items()
          }
          return result

     a__   :      List[Any]     =						dataset_tokenized.map(lowerCamelCase__     ,  batched=lowerCamelCase__     ,  batch_size=1000     ,  num_proc=4      )

     a__   :      Dict     =						0
     a__   :      Tuple     =						0
     for shard in range(0     ,  len(lowerCamelCase__      )     ,  args.shard_size      ):
          a__   :      Dict     =						grouped_dataset[shard : shard + args.shard_size]
          a__   :      Dict     =						len(dataset_snapshot['''input_ids''']      )
          a__   :      Optional[Any]     =						os.path.join(lowerCamelCase__     ,  F'dataset-{shard_count}-{records_containing}.tfrecord'      )
          a__   :      str     =						get_serialized_examples(lowerCamelCase__      )

          with tf.io.TFRecordWriter(lowerCamelCase__      ) as out_file:
               for i in range(len(lowerCamelCase__      )      ):
                    a__   :      Tuple     =						serialized_examples[i]
                    out_file.write(lowerCamelCase__      )
               print('''Wrote file {} containing {} records'''.format(lowerCamelCase__     ,  lowerCamelCase__      )      )

          shard_count += 1
          total_records += records_containing

     with open(F'split-{args.split}-records-count.txt'     ,  '''w'''      ) as f:
          print(F'Total {args.split} records: {total_records}'     ,  file=lowerCamelCase__      )


if __name__ == "__main__":
 __SCREAMING_SNAKE_CASE     =						parse_args()
 main(args)


import math





def      lowerCamelCase_					(							lowerCamelCase__      ,					lowerCamelCase__       ):
				if 0 not in (x, y):
								# We use the relation x^y = y*log10(x), where 10 is the base.
								return y * math.logaa(lowerCamelCase__       )
				else:
								if x == 0:  # 0 raised to any number is 0
												return 0
								elif y == 0:
												return 1  # any number raised to 0 is 1
				raise AssertionError("This should never happen"       )


if __name__ == "__main__":  # Main function
			# Read two numbers from input and typecast them to int using map function.
			# Here x is the base and y is the power.
			__A							='''Enter the base and the power separated by a comma: '''
			__A,		__A							=map(int, input(prompt).split(''','''))
			__A,		__A							=map(int, input(prompt).split(''','''))

			# We find the log of each number, using the function res(), which takes two
			# arguments.
			__A							=res(xa, ya)
			__A							=res(xa, ya)

			# We check for the largest number
			if resa > resa:
						print('''Largest number is''', xa, '''^''', ya)
			elif resa > resa:
						print('''Largest number is''', xa, '''^''', ya)
			else:
						print('''Both are equal''')


