





'''simple docstring'''






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


__lowerCamelCase						: int		=   logging.getLogger(__name__)


def 					_a   ():


  """simple docstring"""

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

  _UpperCamelCase       			=parser.parse_args()
  return args


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  def fn(__SCREAMING_SNAKE_CASE   ):
    return tokenizer(examples['''text''']   )

  return fn


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

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


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=datasets.load_dataset(args.dataset_name  ,      args.dataset_config  ,      split=args.split   )

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

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

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

  _UpperCamelCase       			=dataset_tokenized.map(__SCREAMING_SNAKE_CASE  ,      batched=__SCREAMING_SNAKE_CASE  ,      batch_size=1000  ,      num_proc=4   )

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

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

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


if __name__ == "__main__":
      __lowerCamelCase						: List[str]		=   parse_args()
      main(args)













'''simple docstring'''






from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__lowerCamelCase						: Union[str, Any]		=   logging.get_logger(__name__)

__lowerCamelCase						: int		=   {
    'google/mobilenet_v2_1.4_224': 'https://huggingface.co/google/mobilenet_v2_1.4_224/resolve/main/config.json',
    'google/mobilenet_v2_1.0_224': 'https://huggingface.co/google/mobilenet_v2_1.0_224/resolve/main/config.json',
    'google/mobilenet_v2_0.75_160': 'https://huggingface.co/google/mobilenet_v2_0.75_160/resolve/main/config.json',
    'google/mobilenet_v2_0.35_96': 'https://huggingface.co/google/mobilenet_v2_0.35_96/resolve/main/config.json',
    # See all MobileNetV2 models at https://huggingface.co/models?filter=mobilenet_v2
}







class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				"""mobilenet_v2"""


  def __init__(				self    :					Tuple							,	UpperCamelCase__    :					Union[str, Any]=3							,	UpperCamelCase__    :					Dict=224							,	UpperCamelCase__    :					str=1.0							,	UpperCamelCase__    :					List[Any]=8							,	UpperCamelCase__    :					Union[str, Any]=8							,	UpperCamelCase__    :					str=6							,	UpperCamelCase__    :					str=32							,	UpperCamelCase__    :					str=True							,	UpperCamelCase__    :					int=True							,	UpperCamelCase__    :					str="relu6"							,	UpperCamelCase__    :					Tuple=True							,	UpperCamelCase__    :					List[str]=0.8							,	UpperCamelCase__    :					List[str]=0.02							,	UpperCamelCase__    :					str=0.001							,	UpperCamelCase__    :					Dict=255							,	**UpperCamelCase__    :					Tuple							,	)     ->						List[Any]:
    super().__init__(**UpperCamelCase__  )

    if depth_multiplier <= 0:
      raise ValueError('''depth_multiplier must be greater than zero.'''  )

    _UpperCamelCase       			=num_channels
    _UpperCamelCase       			=image_size
    _UpperCamelCase       			=depth_multiplier
    _UpperCamelCase       			=depth_divisible_by
    _UpperCamelCase       			=min_depth
    _UpperCamelCase       			=expand_ratio
    _UpperCamelCase       			=output_stride
    _UpperCamelCase       			=first_layer_is_expansion
    _UpperCamelCase       			=finegrained_output
    _UpperCamelCase       			=hidden_act
    _UpperCamelCase       			=tf_padding
    _UpperCamelCase       			=classifier_dropout_prob
    _UpperCamelCase       			=initializer_range
    _UpperCamelCase       			=layer_norm_eps
    _UpperCamelCase       			=semantic_loss_ignore_index







class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				version.parse("""1.11""")


  @property
  def   UpperCamelCase__      (				self    :					str  )     ->						Mapping[str, Mapping[int, str]]:
    return OrderedDict([('''pixel_values''', {0: '''batch'''})]  )


  @property
  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Mapping[str, Mapping[int, str]]:
    if self.task == "image-classification":
      return OrderedDict([('''logits''', {0: '''batch'''})]  )
    else:
      return OrderedDict([('''last_hidden_state''', {0: '''batch'''}), ('''pooler_output''', {0: '''batch'''})]  )




  @property
  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						float:
    return 1E-4






