





'''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():



    '''simple docstring'''





    _lowerCAmelCase				=      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."		,       )

    _lowerCAmelCase				=      parser.parse_args()
    return args


def 	__a(SCREAMING_SNAKE_CASE_  :					Tuple    ):



    '''simple docstring'''





    def fn(SCREAMING_SNAKE_CASE_  :					Optional[int]    ):
        return tokenizer(examples["text"]    )

    return fn


def 	__a(SCREAMING_SNAKE_CASE_  :					Dict    ):



    '''simple docstring'''





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


def 	__a(SCREAMING_SNAKE_CASE_  :					Optional[Any]    ):



    '''simple docstring'''





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

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

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

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

    _lowerCAmelCase				=      dataset_tokenized.map(SCREAMING_SNAKE_CASE_		,       batched=SCREAMING_SNAKE_CASE_		,       batch_size=1000		,       num_proc=4    )

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

        with tf.io.TFRecordWriter(SCREAMING_SNAKE_CASE_    ) as out_file:
            for i in range(len(SCREAMING_SNAKE_CASE_    )    ):
                _lowerCAmelCase				=      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__":
       _SCREAMING_SNAKE_CASE	 =							parse_args()
       main(args)









'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_SCREAMING_SNAKE_CASE	 =							{"configuration_xglm": ["XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XGLMConfig"]}

try:
       if not is_sentencepiece_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _SCREAMING_SNAKE_CASE	 =							["XGLMTokenizer"]

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _SCREAMING_SNAKE_CASE	 =							["XGLMTokenizerFast"]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _SCREAMING_SNAKE_CASE	 =							[
           "XGLM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "XGLMForCausalLM",
           "XGLMModel",
           "XGLMPreTrainedModel",
       ]


try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _SCREAMING_SNAKE_CASE	 =							[
           "FlaxXGLMForCausalLM",
           "FlaxXGLMModel",
           "FlaxXGLMPreTrainedModel",
       ]


try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _SCREAMING_SNAKE_CASE	 =							[
           "TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFXGLMForCausalLM",
           "TFXGLMModel",
           "TFXGLMPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_xglm import XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XGLMConfig

       try:
              if not is_sentencepiece_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_xglm import XGLMTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_xglm_fast import XGLMTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_xglm import XGLM_PRETRAINED_MODEL_ARCHIVE_LIST, XGLMForCausalLM, XGLMModel, XGLMPreTrainedModel

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel, FlaxXGLMPreTrainedModel

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_xglm import (
                  TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFXGLMForCausalLM,
                  TFXGLMModel,
                  TFXGLMPreTrainedModel,
              )


else:
       import sys

       _SCREAMING_SNAKE_CASE	 =							_LazyModule(__name__, globals()["__file__"], _import_structure)


