




import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


lowerCamelCase__			:       Union[str, Any]         =							logging.getLogger(__name__)

def  UpperCAmelCase_							(				) ->			str:
      SCREAMING_SNAKE_CASE_     						=      argparse.ArgumentParser(
          description='Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.'							)
      parser.add_argument(
          '--dataset_name'      ,      type=__UpperCAmelCase      ,      default='wikitext'      ,      help='Name of the training. Explore datasets at: hf.co/datasets.'      ,      )
      parser.add_argument(
          '--dataset_config'      ,      type=__UpperCAmelCase      ,      default='wikitext-103-raw-v1'      ,      help='Configuration name of the dataset.'							)
      parser.add_argument(
          '--tokenizer_name_or_path'      ,      type=__UpperCAmelCase      ,      default='sayakpaul/unigram-tokenizer-wikitext'      ,      help='Tokenizer identifier. Can be a local filepath or a Hub identifier.'      ,      )
      parser.add_argument(
          '--shard_size'      ,      type=__UpperCAmelCase      ,      default=10_00      ,      help='Number of entries to go in a single shard.'      ,      )
      parser.add_argument('--split'      ,      type=__UpperCAmelCase      ,      default='train'      ,      choices=['train', 'test', 'validation']							)
      parser.add_argument(
          '--limit'      ,      default=__UpperCAmelCase      ,      type=__UpperCAmelCase      ,      help='Limit the number of shards (used for debugging).'      ,      )
      parser.add_argument(
          '--max_length'      ,      type=__UpperCAmelCase      ,      default=5_12      ,      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=__UpperCAmelCase      ,      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.'      ,      )

      SCREAMING_SNAKE_CASE_     						=      parser.parse_args()
      return args

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Tuple							) ->			Tuple:
      def fn(__UpperCAmelCase						:     Optional[Any]							):
            return tokenizer(examples['text']							)

      return fn

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]							) ->			Optional[int]:
      SCREAMING_SNAKE_CASE_     						=      []
      for i in range(len(tokenized_data['input_ids']							)							):
            SCREAMING_SNAKE_CASE_     						=      {
                '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]							)							),
            }
            SCREAMING_SNAKE_CASE_     						=      tf.train.Features(feature=__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_     						=      tf.train.Example(features=__UpperCAmelCase							)
            SCREAMING_SNAKE_CASE_     						=      example.SerializeToString()
            records.append(__UpperCAmelCase							)
      return records

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Tuple							) ->			List[str]:
      SCREAMING_SNAKE_CASE_     						=      datasets.load_dataset(args.dataset_name      ,      args.dataset_config      ,      split=args.split							)

      if args.limit is not None:
            SCREAMING_SNAKE_CASE_     						=      min(len(__UpperCAmelCase							)      ,      args.limit							)
            SCREAMING_SNAKE_CASE_     						=      dataset.select(range(__UpperCAmelCase							)							)
            print(f"Limiting the dataset to {args.limit} entries."							)

      SCREAMING_SNAKE_CASE_     						=      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							)
            SCREAMING_SNAKE_CASE_     						=      os.path.join(args.output_dir      ,      args.split							)
            if not os.path.exists(__UpperCAmelCase							):
                  os.makedirs(__UpperCAmelCase							)
      else:
            SCREAMING_SNAKE_CASE_     						=      os.path.join(args.output_dir      ,      args.split							)

      # Tokenize the whole dataset at once.
      SCREAMING_SNAKE_CASE_     						=      tokenize_function(__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_     						=      dataset.map(__UpperCAmelCase      ,      batched=__UpperCAmelCase      ,      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(__UpperCAmelCase						:     List[Any]							):
            # Concatenate all texts.
            SCREAMING_SNAKE_CASE_     						=      {k: sum(examples[k]      ,      []							) for k in examples.keys()}
            SCREAMING_SNAKE_CASE_     						=      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 🫀
            SCREAMING_SNAKE_CASE_     						=      (total_length // args.max_length) * args.max_length
            # Split by chunks of max_len.
            SCREAMING_SNAKE_CASE_     						=      {
                k: [t[i : i + args.max_length] for i in range(0      ,      __UpperCAmelCase      ,      args.max_length							)]
                for k, t in concatenated_examples.items()
            }
            return result

      SCREAMING_SNAKE_CASE_     						=      dataset_tokenized.map(__UpperCAmelCase      ,      batched=__UpperCAmelCase      ,      batch_size=10_00      ,      num_proc=4							)

      SCREAMING_SNAKE_CASE_     						=      0
      SCREAMING_SNAKE_CASE_     						=      0
      for shard in range(0      ,      len(__UpperCAmelCase							)      ,      args.shard_size							):
            SCREAMING_SNAKE_CASE_     						=      grouped_dataset[shard : shard + args.shard_size]
            SCREAMING_SNAKE_CASE_     						=      len(dataset_snapshot['input_ids']							)
            SCREAMING_SNAKE_CASE_     						=      os.path.join(__UpperCAmelCase      ,      f"dataset-{shard_count}-{records_containing}.tfrecord"							)
            SCREAMING_SNAKE_CASE_     						=      get_serialized_examples(__UpperCAmelCase							)

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

            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=__UpperCAmelCase							)


if __name__ == "__main__":
   lowerCamelCase__			:       Union[str, Any]         =							parse_args()
   main(args)





from packaging import version

from .import_utils import is_accelerate_available


if is_accelerate_available():
   import accelerate

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]							) ->			int:
      if not is_accelerate_available():
            return method
      SCREAMING_SNAKE_CASE_     						=      version.parse(accelerate.__version__							).base_version
      if version.parse(__UpperCAmelCase							) < version.parse('0.17.0'							):
            return method

      def wrapper(self						:     Optional[int]      ,      *__UpperCAmelCase						:     Optional[Any]      ,      **__UpperCAmelCase						:     Optional[Any]							):
            if hasattr(self      ,      '_hf_hook'							) and hasattr(self._hf_hook      ,      'pre_forward'							):
                  self._hf_hook.pre_forward(self							)
            return method(self      ,      *__UpperCAmelCase      ,      **__UpperCAmelCase							)

      return wrapper