






"""simple docstring"""







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







						"""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 _a			(		_snake_case							):







						"""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 _a			(		_snake_case							):







						"""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 _a			(		_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case				,	_snake_case=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 _a			(		_snake_case							):







						"""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(_snake_case							):
												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(_snake_case							):
												# 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__":
		_UpperCamelCase					      =  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,
)


_UpperCamelCase					      =  {"""configuration_mbart""": ["""MBART_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MBartConfig""", """MBartOnnxConfig"""]}

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

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

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _UpperCamelCase					      =  [
      """MBART_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """MBartForCausalLM""",
      """MBartForConditionalGeneration""",
      """MBartForQuestionAnswering""",
      """MBartForSequenceClassification""",
      """MBartModel""",
      """MBartPreTrainedModel""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _UpperCamelCase					      =  [
      """TFMBartForConditionalGeneration""",
      """TFMBartModel""",
      """TFMBartPreTrainedModel""",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _UpperCamelCase					      =  [
      """FlaxMBartForConditionalGeneration""",
      """FlaxMBartForQuestionAnswering""",
      """FlaxMBartForSequenceClassification""",
      """FlaxMBartModel""",
      """FlaxMBartPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfig, MBartOnnxConfig

  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_mbart import MBartTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_mbart_fast import MBartTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_mbart import (
        MBART_PRETRAINED_MODEL_ARCHIVE_LIST,
        MBartForCausalLM,
        MBartForConditionalGeneration,
        MBartForQuestionAnswering,
        MBartForSequenceClassification,
        MBartModel,
        MBartPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel, TFMBartPreTrainedModel

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_mbart import (
        FlaxMBartForConditionalGeneration,
        FlaxMBartForQuestionAnswering,
        FlaxMBartForSequenceClassification,
        FlaxMBartModel,
        FlaxMBartPreTrainedModel,
    )

else:
  import sys

  _UpperCamelCase					      =  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



