




import argparse
import logging
import os
import re

import tensorflow as tf

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


_lowerCAmelCase						:   Union[str, Any]		=   logging.getLogger(__name__)

_lowerCAmelCase						:   Optional[int]		=   tf.data.AUTOTUNE



def 						lowerCAmelCase		(		):




	"""simple docstring"""







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

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

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

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

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

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

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

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

	UpperCAmelCase__							 =						parser.parse_args()
	return args



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Union[str, Any]	):




	"""simple docstring"""







	try:
		if args.tpu_name:
			UpperCAmelCase__							 =						tf.distribute.cluster_resolver.TPUClusterResolver(
			    args.tpu_name						,							zone=args.tpu_zone						,							project=args.gcp_project	)
		else:
			UpperCAmelCase__							 =						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(_lowerCAmelCase	)
	tf.tpu.experimental.initialize_tpu_system(_lowerCAmelCase	)

	return tpu



def 						lowerCAmelCase		(		_lowerCAmelCase							:   Union[str, Any]	):




	"""simple docstring"""







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

	return num_samples



def 						lowerCAmelCase		(		_lowerCAmelCase							:   str						,							_lowerCAmelCase							:   Optional[int]						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   str						,							_lowerCAmelCase							:   Tuple						,							_lowerCAmelCase							:   List[Any]=None	):




	"""simple docstring"""







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



def 						lowerCAmelCase		(		_lowerCAmelCase							:   str	):




	"""simple docstring"""







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

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

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

	UpperCAmelCase__							 =						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}.'''	)
	UpperCAmelCase__							 =						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}.'''	)

	UpperCAmelCase__							 =						count_samples(_lowerCAmelCase	)

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

	with strategy.scope():
		UpperCAmelCase__							 =						TFAutoModelForMaskedLM.from_config(_lowerCAmelCase	)
		model(model.dummy_inputs	)  # Pass some dummy inputs through the model to ensure all the weights are built
		UpperCAmelCase__ , UpperCAmelCase__							 =						create_optimizer(
		    num_train_steps=_lowerCAmelCase						,							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=_lowerCAmelCase						,							metrics=["accuracy"]	)

	def decode_fn(_lowerCAmelCase							:   Tuple	):
		UpperCAmelCase__							 =						{
		    "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(_lowerCAmelCase						,							_lowerCAmelCase	)

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

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

	UpperCAmelCase__							 =						args.per_replica_batch_size * strategy.num_replicas_in_sync

	UpperCAmelCase__							 =						prepare_dataset(
	    _lowerCAmelCase						,							decode_fn=_lowerCAmelCase						,							mask_fn=_lowerCAmelCase						,							batch_size=_lowerCAmelCase						,							shuffle=_lowerCAmelCase						,							shuffle_buffer_size=args.shuffle_buffer_size						,							)

	UpperCAmelCase__							 =						prepare_dataset(
	    _lowerCAmelCase						,							decode_fn=_lowerCAmelCase						,							mask_fn=_lowerCAmelCase						,							batch_size=_lowerCAmelCase						,							shuffle=_lowerCAmelCase						,							)

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

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

	model.save_pretrained(args.output_dir	)


if __name__ == "__main__":
		_lowerCAmelCase						:   int		=   parse_args()
		main(args)











import warnings

from ...utils import logging
from .image_processing_yolos import YolosImageProcessor


_lowerCAmelCase						:   int		=   logging.get_logger(__name__)
class 			_UpperCamelCase      (				lowerCAmelCase		):
      def __init__(			self   :Optional[Any] ,      *lowerCamelCase   :List[str] ,      **lowerCamelCase   :Tuple      )     ->			None:
       warnings.warn(
           "The class YolosFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
           " use YolosImageProcessor instead." ,      lowerCamelCase ,      )
       super().__init__(*lowerCamelCase ,      **lowerCamelCase      )





