






import argparse
import logging
import os
import re

import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    PushToHubCallback,
    TFAutoModelForMaskedLM,
    create_optimizer,
)


lowerCamelCase      :       Optional[int]    =						logging.getLogger(__name__)

lowerCamelCase      :       Union[str, Any]    =						tf.data.AUTOTUNE





def    __lowerCAmelCase   (    ):
     __lowerCAmelCase             =       argparse.ArgumentParser(description="Train a masked language model on TPU."							)
     parser.add_argument(
         "--pretrained_model_config"       ,			type=__snake_case       ,			default="roberta-base"       ,			help="The model config to use. Note that we don't copy the model's weights, only the config!"       ,			)
     parser.add_argument(
         "--tokenizer"       ,			type=__snake_case       ,			default="unigram-tokenizer-wikitext"       ,			help="The name of the tokenizer to load. We use the pretrained tokenizer to initialize the model's vocab size."       ,			)

     parser.add_argument(
         "--per_replica_batch_size"       ,			type=__snake_case       ,			default=8       ,			help="Batch size per TPU core."       ,			)

     parser.add_argument(
         "--no_tpu"       ,			action="store_true"       ,			help="If set, run on CPU and don't try to initialize a TPU. Useful for debugging on non-TPU instances."       ,			)

     parser.add_argument(
         "--tpu_name"       ,			type=__snake_case       ,			help="Name of TPU resource to initialize. Should be blank on Colab, and 'local' on TPU VMs."       ,			default="local"       ,			)

     parser.add_argument(
         "--tpu_zone"       ,			type=__snake_case       ,			help="Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes."       ,			)

     parser.add_argument(
         "--gcp_project"       ,			type=__snake_case       ,			help="Google cloud project name. Only used for non-Colab TPU nodes."							)

     parser.add_argument(
         "--bfloat16"       ,			action="store_true"       ,			help="Use mixed-precision bfloat16 for training. This is the recommended lower-precision format for TPU."       ,			)

     parser.add_argument(
         "--train_dataset"       ,			type=__snake_case       ,			help="Path to training dataset to load. If the path begins with `gs://`"
         " then the dataset will be loaded from a Google Cloud Storage bucket."       ,			)

     parser.add_argument(
         "--shuffle_buffer_size"       ,			type=__snake_case       ,			default=2**18       ,			help="Size of the shuffle buffer (in samples)"       ,			)

     parser.add_argument(
         "--eval_dataset"       ,			type=__snake_case       ,			help="Path to evaluation dataset to load. If the path begins with `gs://`"
         " then the dataset will be loaded from a Google Cloud Storage bucket."       ,			)

     parser.add_argument(
         "--num_epochs"       ,			type=__snake_case       ,			default=1       ,			help="Number of epochs to train for."       ,			)

     parser.add_argument(
         "--learning_rate"       ,			type=__snake_case       ,			default=1E-4       ,			help="Learning rate to use for training."       ,			)

     parser.add_argument(
         "--weight_decay_rate"       ,			type=__snake_case       ,			default=1E-3       ,			help="Weight decay rate to use for training."       ,			)

     parser.add_argument(
         "--max_length"       ,			type=__snake_case       ,			default=512       ,			help="Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py"       ,			)

     parser.add_argument(
         "--mlm_probability"       ,			type=__snake_case       ,			default=0.15       ,			help="Fraction of tokens to mask during training."       ,			)

     parser.add_argument("--output_dir"       ,			type=__snake_case       ,			required=__snake_case       ,			help="Path to save model checkpoints to."							)
     parser.add_argument("--hub_model_id"       ,			type=__snake_case       ,			help="Model ID to upload to on the Hugging Face Hub."							)

     __lowerCAmelCase             =       parser.parse_args()
     return args





def    __lowerCAmelCase   (    __snake_case							):
     try:
          if args.tpu_name:
               __lowerCAmelCase             =       tf.distribute.cluster_resolver.TPUClusterResolver(
                   args.tpu_name       ,			zone=args.tpu_zone       ,			project=args.gcp_project							)
          else:
               __lowerCAmelCase             =       tf.distribute.cluster_resolver.TPUClusterResolver()
     except ValueError:
          raise RuntimeError(
              "Couldn't connect to TPU! Most likely you need to specify --tpu_name, --tpu_zone, or "
              "--gcp_project. When running on a TPU VM, use --tpu_name local."							)

     tf.config.experimental_connect_to_cluster(__snake_case							)
     tf.tpu.experimental.initialize_tpu_system(__snake_case							)

     return tpu





def    __lowerCAmelCase   (    __snake_case							):
     __lowerCAmelCase             =       0
     for file in file_list:
          __lowerCAmelCase             =       file.split("/"							)[-1]
          __lowerCAmelCase             =       re.search(r"-\d+-(\d+)\.tfrecord"       ,			__snake_case							).group(1							)
          __lowerCAmelCase             =       int(__snake_case							)
          num_samples += sample_count

     return num_samples





def    __lowerCAmelCase   (    __snake_case       ,			__snake_case       ,			__snake_case       ,			__snake_case       ,			__snake_case       ,			__snake_case=None							):
     __lowerCAmelCase             =       count_samples(__snake_case							)
     __lowerCAmelCase             =       tf.data.Dataset.from_tensor_slices(__snake_case							)
     if shuffle:
          __lowerCAmelCase             =       dataset.shuffle(len(__snake_case							)							)
     __lowerCAmelCase             =       tf.data.TFRecordDataset(__snake_case       ,			num_parallel_reads=__snake_case							)
     # TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
     __lowerCAmelCase             =       dataset.apply(tf.data.experimental.assert_cardinality(__snake_case							)							)
     __lowerCAmelCase             =       dataset.map(__snake_case       ,			num_parallel_calls=__snake_case							)
     if shuffle:
          assert shuffle_buffer_size is not None
          __lowerCAmelCase             =       dataset.shuffle(args.shuffle_buffer_size							)
     __lowerCAmelCase             =       dataset.batch(__snake_case       ,			drop_remainder=__snake_case							)
     __lowerCAmelCase             =       dataset.map(__snake_case       ,			num_parallel_calls=__snake_case							)
     __lowerCAmelCase             =       dataset.prefetch(__snake_case							)
     return dataset





def    __lowerCAmelCase   (    __snake_case							):
     if not args.no_tpu:
          __lowerCAmelCase             =       initialize_tpu(__snake_case							)
          __lowerCAmelCase             =       tf.distribute.TPUStrategy(__snake_case							)
     else:
          __lowerCAmelCase             =       tf.distribute.OneDeviceStrategy(device="/gpu:0"							)

     if args.bfloataa:
          tf.keras.mixed_precision.set_global_policy("mixed_bfloat16"							)

     __lowerCAmelCase             =       AutoTokenizer.from_pretrained(args.tokenizer							)
     __lowerCAmelCase             =       AutoConfig.from_pretrained(args.pretrained_model_config							)
     __lowerCAmelCase             =       tokenizer.vocab_size

     __lowerCAmelCase             =       tf.io.gfile.glob(os.path.join(args.train_dataset       ,			"*.tfrecord"							)							)
     if not training_records:
          raise ValueError(F"""No .tfrecord files found in {args.train_dataset}."""							)
     __lowerCAmelCase             =       tf.io.gfile.glob(os.path.join(args.eval_dataset       ,			"*.tfrecord"							)							)
     if not eval_records:
          raise ValueError(F"""No .tfrecord files found in {args.eval_dataset}."""							)

     __lowerCAmelCase             =       count_samples(__snake_case							)

     __lowerCAmelCase             =       num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
     __lowerCAmelCase             =       steps_per_epoch * args.num_epochs

     with strategy.scope():
          __lowerCAmelCase             =       TFAutoModelForMaskedLM.from_config(__snake_case							)
          model(model.dummy_inputs							)  # Pass some dummy inputs through the model to ensure all the weights are built
          __lowerCAmelCase      ,			__lowerCAmelCase             =       create_optimizer(
              num_train_steps=__snake_case       ,			num_warmup_steps=total_train_steps // 20       ,			init_lr=args.learning_rate       ,			weight_decay_rate=args.weight_decay_rate       ,			)

          # Transformers models compute the right loss for their task by default when labels are passed, and will
          # use this for training unless you specify your own loss function in compile().
          model.compile(optimizer=__snake_case       ,			metrics=["accuracy"]							)

     def decode_fn(__snake_case							):
          __lowerCAmelCase             =       {
              "input_ids": tf.io.FixedLenFeature(dtype=tf.intaa       ,			shape=(args.max_length,)							),
              "attention_mask": tf.io.FixedLenFeature(dtype=tf.intaa       ,			shape=(args.max_length,)							),
          }
          return tf.io.parse_single_example(__snake_case       ,			__snake_case							)

     # Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
     # use their methods in our data pipeline.
     __lowerCAmelCase             =       DataCollatorForLanguageModeling(
         tokenizer=__snake_case       ,			mlm_probability=args.mlm_probability       ,			mlm=__snake_case       ,			return_tensors="tf"							)

     def mask_with_collator(__snake_case							):
          # TF really needs an isin() function
          __lowerCAmelCase             =       (
              ~tf.cast(batch["attention_mask"]       ,			tf.bool							)
              | (batch["input_ids"] == tokenizer.cls_token_id)
              | (batch["input_ids"] == tokenizer.sep_token_id)
          )
          __lowerCAmelCase      ,			__lowerCAmelCase             =       data_collator.tf_mask_tokens(
              batch["input_ids"]       ,			vocab_size=len(__snake_case							)       ,			mask_token_id=tokenizer.mask_token_id       ,			special_tokens_mask=__snake_case       ,			)
          return batch

     __lowerCAmelCase             =       args.per_replica_batch_size * strategy.num_replicas_in_sync

     __lowerCAmelCase             =       prepare_dataset(
         __snake_case       ,			decode_fn=__snake_case       ,			mask_fn=__snake_case       ,			batch_size=__snake_case       ,			shuffle=__snake_case       ,			shuffle_buffer_size=args.shuffle_buffer_size       ,			)

     __lowerCAmelCase             =       prepare_dataset(
         __snake_case       ,			decode_fn=__snake_case       ,			mask_fn=__snake_case       ,			batch_size=__snake_case       ,			shuffle=__snake_case       ,			)

     __lowerCAmelCase             =       []
     if args.hub_model_id:
          callbacks.append(
              PushToHubCallback(output_dir=args.output_dir       ,			hub_model_id=args.hub_model_id       ,			tokenizer=__snake_case							)							)

     model.fit(
         __snake_case       ,			validation_data=__snake_case       ,			epochs=args.num_epochs       ,			callbacks=__snake_case       ,			)

     model.save_pretrained(args.output_dir							)


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














import warnings
from typing import List, Optional, Union

from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType



class 				_UpperCamelCase      (a_      ):
      snake_case_	=				["""image_processor""", """tokenizer"""]
      snake_case_	=				"""FlavaImageProcessor"""
      snake_case_	=				("""BertTokenizer""", """BertTokenizerFast""")


      def __init__(						self       ,    __UpperCamelCase=None       ,    __UpperCamelCase=None       ,    **__UpperCamelCase						)->						Tuple:
           __lowerCAmelCase             =       None
           if "feature_extractor" in kwargs:
                warnings.warn(
                    "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
                    " instead."       ,    __UpperCamelCase       ,    )
                __lowerCAmelCase             =       kwargs.pop("feature_extractor"						)

           __lowerCAmelCase             =       image_processor if image_processor is not None else feature_extractor
           if image_processor is None:
                raise ValueError("You need to specify an `image_processor`."						)
           if tokenizer is None:
                raise ValueError("You need to specify a `tokenizer`."						)

           super().__init__(__UpperCamelCase       ,    __UpperCamelCase						)
           __lowerCAmelCase             =       self.image_processor


      def __call__(						self       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = True       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = None       ,    __UpperCamelCase = 0       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = False       ,    __UpperCamelCase = True       ,    __UpperCamelCase = None       ,    **__UpperCamelCase       ,    )->						int:

           if text is None and images is None:
                raise ValueError("You have to specify either text or images. Both cannot be none."						)

           if text is not None:
                __lowerCAmelCase             =       self.tokenizer(
                    text=__UpperCamelCase       ,    add_special_tokens=__UpperCamelCase       ,    padding=__UpperCamelCase       ,    truncation=__UpperCamelCase       ,    max_length=__UpperCamelCase       ,    stride=__UpperCamelCase       ,    pad_to_multiple_of=__UpperCamelCase       ,    return_token_type_ids=__UpperCamelCase       ,    return_attention_mask=__UpperCamelCase       ,    return_overflowing_tokens=__UpperCamelCase       ,    return_special_tokens_mask=__UpperCamelCase       ,    return_offsets_mapping=__UpperCamelCase       ,    return_length=__UpperCamelCase       ,    verbose=__UpperCamelCase       ,    return_tensors=__UpperCamelCase       ,    **__UpperCamelCase       ,    )
           if images is not None:
                __lowerCAmelCase             =       self.image_processor(
                    __UpperCamelCase       ,    return_image_mask=__UpperCamelCase       ,    return_codebook_pixels=__UpperCamelCase       ,    return_tensors=__UpperCamelCase       ,    **__UpperCamelCase       ,    )

           if text is not None and images is not None:
                encoding.update(__UpperCamelCase						)
                return encoding
           elif text is not None:
                return encoding
           else:
                return BatchEncoding(data=dict(**__UpperCamelCase						)       ,    tensor_type=__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self       ,    *__UpperCamelCase       ,    **__UpperCamelCase						)->						int:
           return self.tokenizer.batch_decode(*__UpperCamelCase       ,    **__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self       ,    *__UpperCamelCase       ,    **__UpperCamelCase						)->						Optional[int]:
           return self.tokenizer.decode(*__UpperCamelCase       ,    **__UpperCamelCase						)


      @property
      def 			__UpperCAmelCase		(						self						)->						Any:
           __lowerCAmelCase             =       self.tokenizer.model_input_names
           __lowerCAmelCase             =       self.image_processor.model_input_names
           return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names						)						)


      @property
      def 			__UpperCAmelCase		(						self						)->						Dict:
           warnings.warn(
               "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."       ,    __UpperCamelCase       ,    )
           return self.image_processor_class


      @property
      def 			__UpperCAmelCase		(						self						)->						Union[str, Any]:
           warnings.warn(
               "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."       ,    __UpperCamelCase       ,    )
           return self.image_processor






