






"""simple docstring"""
import argparse
import logging
import os
import re

import tensorflow as tf

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


lowerCamelCase_       			=   logging.getLogger(__name__)

lowerCamelCase_       			=   tf.data.AUTOTUNE

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

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

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

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

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

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

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

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

			UpperCAmelCase_  :				Union[str, Any]      =      parser.parse_args()
			return args

def        snake_case   (     A__		):
			try:
						if args.tpu_name:
									UpperCAmelCase_  :				int      =      tf.distribute.cluster_resolver.TPUClusterResolver(
									    args.tpu_name			,zone=args.tpu_zone			,project=args.gcp_project		)
						else:
									UpperCAmelCase_  :				List[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(A__		)
			tf.tpu.experimental.initialize_tpu_system(A__		)

			return tpu

def        snake_case   (     A__		):
			UpperCAmelCase_  :				Optional[Any]      =      0
			for file in file_list:
						UpperCAmelCase_  :				List[str]      =      file.split("/"		)[-1]
						UpperCAmelCase_  :				int      =      re.search(r"-\d+-(\d+)\.tfrecord"			,A__		).group(1		)
						UpperCAmelCase_  :				List[str]      =      int(A__		)
						num_samples += sample_count

			return num_samples

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

def        snake_case   (     A__		):
			if not args.no_tpu:
						UpperCAmelCase_  :				Union[str, Any]      =      initialize_tpu(A__		)
						UpperCAmelCase_  :				List[str]      =      tf.distribute.TPUStrategy(A__		)
			else:
						UpperCAmelCase_  :				List[Any]      =      tf.distribute.OneDeviceStrategy(device="/gpu:0"		)

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

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

			UpperCAmelCase_  :				Any      =      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_  :				int      =      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_  :				Dict      =      count_samples(A__		)

			UpperCAmelCase_  :				List[str]      =      num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
			UpperCAmelCase_  :				Optional[Any]      =      steps_per_epoch * args.num_epochs

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

			def decode_fn(A__		):
						UpperCAmelCase_  :				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(A__			,A__		)

			# 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_  :				Tuple      =      DataCollatorForLanguageModeling(
			    tokenizer=A__			,mlm_probability=args.mlm_probability			,mlm=A__			,return_tensors="tf"		)

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

			UpperCAmelCase_  :				Optional[int]      =      args.per_replica_batch_size * strategy.num_replicas_in_sync

			UpperCAmelCase_  :				Tuple      =      prepare_dataset(
			    A__			,decode_fn=A__			,mask_fn=A__			,batch_size=A__			,shuffle=A__			,shuffle_buffer_size=args.shuffle_buffer_size			,)

			UpperCAmelCase_  :				Union[str, Any]      =      prepare_dataset(
			    A__			,decode_fn=A__			,mask_fn=A__			,batch_size=A__			,shuffle=A__			,)

			UpperCAmelCase_  :				Optional[int]      =      []
			if args.hub_model_id:
						callbacks.append(
						    PushToHubCallback(output_dir=args.output_dir			,hub_model_id=args.hub_model_id			,tokenizer=A__		)		)

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

			model.save_pretrained(args.output_dir		)


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











"""simple docstring"""
from typing import TYPE_CHECKING

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


lowerCamelCase_       			=   {
    '''configuration_funnel''': ['''FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''FunnelConfig'''],
    '''convert_funnel_original_tf_checkpoint_to_pytorch''': [],
    '''tokenization_funnel''': ['''FunnelTokenizer'''],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   ['''FunnelTokenizerFast''']

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''FunnelBaseModel''',
			    '''FunnelForMaskedLM''',
			    '''FunnelForMultipleChoice''',
			    '''FunnelForPreTraining''',
			    '''FunnelForQuestionAnswering''',
			    '''FunnelForSequenceClassification''',
			    '''FunnelForTokenClassification''',
			    '''FunnelModel''',
			    '''FunnelPreTrainedModel''',
			    '''load_tf_weights_in_funnel''',
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCamelCase_       			=   [
			    '''TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''TFFunnelBaseModel''',
			    '''TFFunnelForMaskedLM''',
			    '''TFFunnelForMultipleChoice''',
			    '''TFFunnelForPreTraining''',
			    '''TFFunnelForQuestionAnswering''',
			    '''TFFunnelForSequenceClassification''',
			    '''TFFunnelForTokenClassification''',
			    '''TFFunnelModel''',
			    '''TFFunnelPreTrainedModel''',
			]


if TYPE_CHECKING:
			from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
			from .tokenization_funnel import FunnelTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_funnel_fast import FunnelTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_funnel import (
						    FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
						    FunnelBaseModel,
						    FunnelForMaskedLM,
						    FunnelForMultipleChoice,
						    FunnelForPreTraining,
						    FunnelForQuestionAnswering,
						    FunnelForSequenceClassification,
						    FunnelForTokenClassification,
						    FunnelModel,
						    FunnelPreTrainedModel,
						    load_tf_weights_in_funnel,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_funnel import (
						    TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFFunnelBaseModel,
						    TFFunnelForMaskedLM,
						    TFFunnelForMultipleChoice,
						    TFFunnelForPreTraining,
						    TFFunnelForQuestionAnswering,
						    TFFunnelForSequenceClassification,
						    TFFunnelForTokenClassification,
						    TFFunnelModel,
						    TFFunnelPreTrainedModel,
						)

else:
			import sys

			lowerCamelCase_       			=   _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



