

import argparse
import logging
import os
import re

import tensorflow as tf

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


_UpperCamelCase 			=		logging.getLogger(__name__)

_UpperCamelCase 			=		tf.data.AUTOTUNE







def 		_lowercase (     ):
      __lowerCAmelCase					:						Dict					=					argparse.ArgumentParser(description='''Train a masked language model on TPU.'''						)
      parser.add_argument(
          '''--pretrained_model_config'''   ,	type=lowercase__   ,	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=lowercase__   ,	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=lowercase__   ,	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=lowercase__   ,	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=lowercase__   ,	help='''Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes.'''   ,	)

      parser.add_argument(
          '''--gcp_project'''   ,	type=lowercase__   ,	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=lowercase__   ,	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=lowercase__   ,	default=2**1_8   ,	help='''Size of the shuffle buffer (in samples)'''   ,	)

      parser.add_argument(
          '''--eval_dataset'''   ,	type=lowercase__   ,	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=lowercase__   ,	default=1   ,	help='''Number of epochs to train for.'''   ,	)

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

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

      parser.add_argument(
          '''--max_length'''   ,	type=lowercase__   ,	default=5_1_2   ,	help='''Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py'''   ,	)

      parser.add_argument(
          '''--mlm_probability'''   ,	type=lowercase__   ,	default=0.1_5   ,	help='''Fraction of tokens to mask during training.'''   ,	)

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

      __lowerCAmelCase					:						Any					=					parser.parse_args()
      return args







def 		_lowercase (     lowercase__						):
      try:
            if args.tpu_name:
                  __lowerCAmelCase					:						Union[str, Any]					=					tf.distribute.cluster_resolver.TPUClusterResolver(
                      args.tpu_name   ,	zone=args.tpu_zone   ,	project=args.gcp_project						)
            else:
                  __lowerCAmelCase					:						str					=					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(lowercase__						)
      tf.tpu.experimental.initialize_tpu_system(lowercase__						)

      return tpu







def 		_lowercase (     lowercase__						):
      __lowerCAmelCase					:						List[Any]					=					0
      for file in file_list:
            __lowerCAmelCase					:						int					=					file.split('''/'''						)[-1]
            __lowerCAmelCase					:						List[Any]					=					re.search(r'''-\d+-(\d+)\.tfrecord'''   ,	lowercase__						).group(1						)
            __lowerCAmelCase					:						str					=					int(lowercase__						)
            num_samples += sample_count

      return num_samples







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







def 		_lowercase (     lowercase__						):
      if not args.no_tpu:
            __lowerCAmelCase					:						Optional[Any]					=					initialize_tpu(lowercase__						)
            __lowerCAmelCase					:						Optional[Any]					=					tf.distribute.TPUStrategy(lowercase__						)
      else:
            __lowerCAmelCase					:						List[str]					=					tf.distribute.OneDeviceStrategy(device='''/gpu:0'''						)

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

      __lowerCAmelCase					:						str					=					AutoTokenizer.from_pretrained(args.tokenizer						)
      __lowerCAmelCase					:						str					=					AutoConfig.from_pretrained(args.pretrained_model_config						)
      __lowerCAmelCase					:						List[Any]					=					tokenizer.vocab_size

      __lowerCAmelCase					:						str					=					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					:						str					=					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					:						Dict					=					count_samples(lowercase__						)

      __lowerCAmelCase					:						Optional[Any]					=					num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
      __lowerCAmelCase					:						Union[str, Any]					=					steps_per_epoch * args.num_epochs

      with strategy.scope():
            __lowerCAmelCase					:						Optional[Any]					=					TFAutoModelForMaskedLM.from_config(lowercase__						)
            model(model.dummy_inputs						)  # Pass some dummy inputs through the model to ensure all the weights are built
            __lowerCAmelCase,							__lowerCAmelCase					:						Optional[Any]					=					create_optimizer(
                num_train_steps=lowercase__   ,	num_warmup_steps=total_train_steps // 2_0   ,	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=lowercase__   ,	metrics=['''accuracy''']						)

      def decode_fn(lowercase__						):
            __lowerCAmelCase					:						Optional[int]					=					{
                '''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(lowercase__   ,	lowercase__						)

      # 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					:						str					=					DataCollatorForLanguageModeling(
          tokenizer=lowercase__   ,	mlm_probability=args.mlm_probability   ,	mlm=lowercase__   ,	return_tensors='''tf'''						)

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

      __lowerCAmelCase					:						str					=					args.per_replica_batch_size * strategy.num_replicas_in_sync

      __lowerCAmelCase					:						Tuple					=					prepare_dataset(
          lowercase__   ,	decode_fn=lowercase__   ,	mask_fn=lowercase__   ,	batch_size=lowercase__   ,	shuffle=lowercase__   ,	shuffle_buffer_size=args.shuffle_buffer_size   ,	)

      __lowerCAmelCase					:						Union[str, Any]					=					prepare_dataset(
          lowercase__   ,	decode_fn=lowercase__   ,	mask_fn=lowercase__   ,	batch_size=lowercase__   ,	shuffle=lowercase__   ,	)

      __lowerCAmelCase					:						Union[str, Any]					=					[]
      if args.hub_model_id:
            callbacks.append(
                PushToHubCallback(output_dir=args.output_dir   ,	hub_model_id=args.hub_model_id   ,	tokenizer=lowercase__						)						)

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

      model.save_pretrained(args.output_dir						)


if __name__ == "__main__":
 _UpperCamelCase 			=		parse_args()
 main(args)





import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
 import torch

if is_vision_available():
 import PIL


_UpperCamelCase 			=		logging.get_logger(__name__)







def 		_lowercase (     lowercase__   ,	lowercase__   ,	lowercase__   ,	lowercase__						):
      def constraint_to_multiple_of(lowercase__   ,	lowercase__   ,	lowercase__=0   ,	lowercase__=None						):
            __lowerCAmelCase					:						int					=					round(val / multiple						) * multiple

            if max_val is not None and x > max_val:
                  __lowerCAmelCase					:						Optional[int]					=					math.floor(val / multiple						) * multiple

            if x < min_val:
                  __lowerCAmelCase					:						Any					=					math.ceil(val / multiple						) * multiple

            return x

      __lowerCAmelCase					:						Dict					=					(output_size, output_size) if isinstance(lowercase__   ,	lowercase__						) else output_size

      __lowerCAmelCase,							__lowerCAmelCase					:						Optional[Any]					=					get_image_size(lowercase__						)
      __lowerCAmelCase,							__lowerCAmelCase					:						int					=					output_size

      # determine new height and width
      __lowerCAmelCase					:						Optional[Any]					=					output_height / input_height
      __lowerCAmelCase					:						List[Any]					=					output_width / input_width

      if keep_aspect_ratio:
            # scale as little as possible
            if abs(1 - scale_width						) < abs(1 - scale_height						):
                  # fit width
                  __lowerCAmelCase					:						str					=					scale_width
            else:
                  # fit height
                  __lowerCAmelCase					:						str					=					scale_height

      __lowerCAmelCase					:						Any					=					constraint_to_multiple_of(scale_height * input_height   ,	multiple=lowercase__						)
      __lowerCAmelCase					:						Union[str, Any]					=					constraint_to_multiple_of(scale_width * input_width   ,	multiple=lowercase__						)

      return (new_height, new_width)







class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  ["""pixel_values"""]
 def __init__(							self      ,							A_ = True      ,							A_ = None      ,							A_ = PILImageResampling.BILINEAR      ,							A_ = False      ,							A_ = 1      ,							A_ = True      ,							A_ = 1 / 255      ,							A_ = True      ,							A_ = None      ,							A_ = None      ,							**A_      ,							)					->None:



       '''simple docstring'''





       super().__init__(**A_					)
       __lowerCAmelCase					:						Union[str, Any]					=					size if size is not None else {'''height''': 384, '''width''': 384}
       __lowerCAmelCase					:						Dict					=					get_size_dict(A_					)
       __lowerCAmelCase					:						Optional[Any]					=					do_resize
       __lowerCAmelCase					:						int					=					size
       __lowerCAmelCase					:						Dict					=					keep_aspect_ratio
       __lowerCAmelCase					:						List[Any]					=					ensure_multiple_of
       __lowerCAmelCase					:						Tuple					=					resample
       __lowerCAmelCase					:						Dict					=					do_rescale
       __lowerCAmelCase					:						Any					=					rescale_factor
       __lowerCAmelCase					:						List[Any]					=					do_normalize
       __lowerCAmelCase					:						Optional[int]					=					image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
       __lowerCAmelCase					:						Optional[int]					=					image_std if image_std is not None else IMAGENET_STANDARD_STD
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_ = False      ,							A_ = 1      ,							A_ = PILImageResampling.BICUBIC      ,							A_ = None      ,							**A_      ,							)					->np.ndarray:



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					get_size_dict(A_					)
       if "height" not in size or "width" not in size:
             raise ValueError(f"""The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"""					)
       __lowerCAmelCase					:						Union[str, Any]					=					get_resize_output_image_size(
           A_      ,							output_size=(size['''height'''], size['''width'''])      ,							keep_aspect_ratio=A_      ,							multiple=A_      ,							)
       return resize(A_      ,							size=A_      ,							resample=A_      ,							data_format=A_      ,							**A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_ = None      ,							**A_      ,							)					->Dict:



       '''simple docstring'''





       return rescale(A_      ,							scale=A_      ,							data_format=A_      ,							**A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_ = None      ,							**A_      ,							)					->np.ndarray:



       '''simple docstring'''





       return normalize(A_      ,							mean=A_      ,							std=A_      ,							data_format=A_      ,							**A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = None      ,							A_ = ChannelDimension.FIRST      ,							**A_      ,							)					->PIL.Image.Image:



       '''simple docstring'''





       __lowerCAmelCase					:						int					=					do_resize if do_resize is not None else self.do_resize
       __lowerCAmelCase					:						Optional[int]					=					size if size is not None else self.size
       __lowerCAmelCase					:						Union[str, Any]					=					get_size_dict(A_					)
       __lowerCAmelCase					:						List[Any]					=					keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
       __lowerCAmelCase					:						Optional[int]					=					ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
       __lowerCAmelCase					:						Tuple					=					resample if resample is not None else self.resample
       __lowerCAmelCase					:						Optional[int]					=					do_rescale if do_rescale is not None else self.do_rescale
       __lowerCAmelCase					:						Union[str, Any]					=					rescale_factor if rescale_factor is not None else self.rescale_factor
       __lowerCAmelCase					:						Tuple					=					do_normalize if do_normalize is not None else self.do_normalize
       __lowerCAmelCase					:						str					=					image_mean if image_mean is not None else self.image_mean
       __lowerCAmelCase					:						Optional[Any]					=					image_std if image_std is not None else self.image_std

       __lowerCAmelCase					:						Optional[Any]					=					make_list_of_images(A_					)

       if not valid_images(A_					):
             raise ValueError(
                 '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                 '''torch.Tensor, tf.Tensor or jax.ndarray.'''					)

       if do_resize and size is None or resample is None:
             raise ValueError('''Size and resample must be specified if do_resize is True.'''					)

       if do_rescale and rescale_factor is None:
             raise ValueError('''Rescale factor must be specified if do_rescale is True.'''					)

       if do_normalize and (image_mean is None or image_std is None):
             raise ValueError('''Image mean and std must be specified if do_normalize is True.'''					)

       # All transformations expect numpy arrays.
       __lowerCAmelCase					:						Any					=					[to_numpy_array(A_					) for image in images]

       if do_resize:
             __lowerCAmelCase					:						Optional[Any]					=					[self.resize(image=A_      ,							size=A_      ,							resample=A_					) for image in images]

       if do_rescale:
             __lowerCAmelCase					:						Tuple					=					[self.rescale(image=A_      ,							scale=A_					) for image in images]

       if do_normalize:
             __lowerCAmelCase					:						str					=					[self.normalize(image=A_      ,							mean=A_      ,							std=A_					) for image in images]

       __lowerCAmelCase					:						Union[str, Any]					=					[to_channel_dimension_format(A_      ,							A_					) for image in images]

       __lowerCAmelCase					:						Dict					=					{'''pixel_values''': images}
       return BatchFeature(data=A_      ,							tensor_type=A_					)





 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None					)					->Any:



       '''simple docstring'''





       __lowerCAmelCase					:						Any					=					outputs.logits

       # Resize logits and compute semantic segmentation maps
       if target_sizes is not None:
             if len(A_					) != len(A_					):
                   raise ValueError(
                       '''Make sure that you pass in as many target sizes as the batch dimension of the logits'''					)

             if is_torch_tensor(A_					):
                   __lowerCAmelCase					:						Optional[int]					=					target_sizes.numpy()

             __lowerCAmelCase					:						List[str]					=					[]

             for idx in range(len(A_					)					):
                   __lowerCAmelCase					:						Any					=					torch.nn.functional.interpolate(
                       logits[idx].unsqueeze(dim=0					)      ,							size=target_sizes[idx]      ,							mode='''bilinear'''      ,							align_corners=A_					)
                   __lowerCAmelCase					:						str					=					resized_logits[0].argmax(dim=0					)
                   semantic_segmentation.append(A_					)
       else:
             __lowerCAmelCase					:						Any					=					logits.argmax(dim=1					)
             __lowerCAmelCase					:						List[Any]					=					[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]					)]

       return semantic_segmentation


