






"""simple docstring"""




import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


_snake_case          =   logging.getLogger(__name__)

def 			lowerCAmelCase__     (     ):



  '''simple docstring'''
  _a  :       str             =	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=1_0_0_0     ,   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_1_2     ,   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."""     ,   )

  _a  :       List[Any]             =	parser.parse_args()
  return args

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  def fn(UpperCamelCase__				):
    return tokenizer(examples["""text"""]				)

  return fn

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       List[Any]             =	[]
  for i in range(len(tokenized_data["""input_ids"""]				)				):
    _a  :       Optional[int]             =	{
        """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  :       Any             =	tf.train.Features(feature=UpperCamelCase__				)
    _a  :       Tuple             =	tf.train.Example(features=UpperCamelCase__				)
    _a  :       Dict             =	example.SerializeToString()
    records.append(UpperCamelCase__				)
  return records

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       List[Any]             =	datasets.load_dataset(args.dataset_name     ,   args.dataset_config     ,   split=args.split				)

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

  _a  :       List[Any]             =	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  :       str             =	os.path.join(args.output_dir     ,   args.split				)
    if not os.path.exists(UpperCamelCase__				):
      os.makedirs(UpperCamelCase__				)
  else:
    _a  :       Union[str, Any]             =	os.path.join(args.output_dir     ,   args.split				)

  # Tokenize the whole dataset at once.
  _a  :       Optional[Any]             =	tokenize_function(UpperCamelCase__				)
  _a  :       Tuple             =	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__				):
    # Concatenate all texts.
    _a  :       Union[str, Any]             =	{k: sum(examples[k]     ,   []				) for k in examples.keys()}
    _a  :       Any             =	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  :       Tuple             =	(total_length // args.max_length) * args.max_length
    # Split by chunks of max_len.
    _a  :       List[str]             =	{
        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

  _a  :       int             =	dataset_tokenized.map(UpperCamelCase__     ,   batched=UpperCamelCase__     ,   batch_size=1_0_0_0     ,   num_proc=4				)

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

    with tf.io.TFRecordWriter(UpperCamelCase__				) as out_file:
      for i in range(len(UpperCamelCase__				)				):
        _a  :       Optional[Any]             =	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__":
 _snake_case          =   parse_args()
 main(args)









"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


_snake_case          =   {
    'configuration_transfo_xl': ['TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TransfoXLConfig'],
    'tokenization_transfo_xl': ['TransfoXLCorpus', 'TransfoXLTokenizer'],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case          =   [
     'TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'AdaptiveEmbedding',
     'TransfoXLForSequenceClassification',
     'TransfoXLLMHeadModel',
     'TransfoXLModel',
     'TransfoXLPreTrainedModel',
     'load_tf_weights_in_transfo_xl',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case          =   [
     'TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST',
     'TFAdaptiveEmbedding',
     'TFTransfoXLForSequenceClassification',
     'TFTransfoXLLMHeadModel',
     'TFTransfoXLMainLayer',
     'TFTransfoXLModel',
     'TFTransfoXLPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
 from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_transfo_xl import (
      TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
      AdaptiveEmbedding,
      TransfoXLForSequenceClassification,
      TransfoXLLMHeadModel,
      TransfoXLModel,
      TransfoXLPreTrainedModel,
      load_tf_weights_in_transfo_xl,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_transfo_xl import (
      TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAdaptiveEmbedding,
      TFTransfoXLForSequenceClassification,
      TFTransfoXLLMHeadModel,
      TFTransfoXLMainLayer,
      TFTransfoXLModel,
      TFTransfoXLPreTrainedModel,
  )

else:
 import sys

 _snake_case          =   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

