



import argparse
import logging
import os
import re

import tensorflow as tf

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


A__								=						logging.getLogger(__name__)

A__								=						tf.data.AUTOTUNE



def 			_lowerCAmelCase       (					)   ->					List[Any]:


				"""simple docstring"""

				snake_case__	:      Tuple     		=	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.'''      )

				snake_case__	:      List[str]     		=	parser.parse_args()
				return args



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					Dict:


				"""simple docstring"""

				try:
								if args.tpu_name:
												snake_case__	:      Optional[int]     		=	tf.distribute.cluster_resolver.TPUClusterResolver(
												    args.tpu_name  ,  zone=args.tpu_zone  ,  project=args.gcp_project      )
								else:
												snake_case__	:      List[Any]     		=	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      )   ->					str:


				"""simple docstring"""

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

				return num_samples



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase=None      )   ->					Any:


				"""simple docstring"""

				snake_case__	:      List[Any]     		=	count_samples(__lowerCAmelCase      )
				snake_case__	:      List[str]     		=	tf.data.Dataset.from_tensor_slices(__lowerCAmelCase      )
				if shuffle:
								snake_case__	:      Dict     		=	dataset.shuffle(len(__lowerCAmelCase      )      )
				snake_case__	:      Union[str, Any]     		=	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
				snake_case__	:      Dict     		=	dataset.apply(tf.data.experimental.assert_cardinality(__lowerCAmelCase      )      )
				snake_case__	:      Optional[Any]     		=	dataset.map(__lowerCAmelCase  ,  num_parallel_calls=__lowerCAmelCase      )
				if shuffle:
								assert shuffle_buffer_size is not None
								snake_case__	:      Dict     		=	dataset.shuffle(args.shuffle_buffer_size      )
				snake_case__	:      Optional[int]     		=	dataset.batch(__lowerCAmelCase  ,  drop_remainder=__lowerCAmelCase      )
				snake_case__	:      Optional[Any]     		=	dataset.map(__lowerCAmelCase  ,  num_parallel_calls=__lowerCAmelCase      )
				snake_case__	:      List[str]     		=	dataset.prefetch(__lowerCAmelCase      )
				return dataset



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					Optional[Any]:


				"""simple docstring"""

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

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

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

				snake_case__	:      Union[str, 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}."""      )
				snake_case__	:      List[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}."""      )

				snake_case__	:      Union[str, Any]     		=	count_samples(__lowerCAmelCase      )

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

				with strategy.scope():
								snake_case__	:      str     		=	TFAutoModelForMaskedLM.from_config(__lowerCAmelCase      )
								model(model.dummy_inputs      )  # Pass some dummy inputs through the model to ensure all the weights are built
								snake_case__		, snake_case__	:      List[Any]     		=	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      ):
								snake_case__	:      Optional[Any]     		=	{
								    '''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.
				snake_case__	:      Tuple     		=	DataCollatorForLanguageModeling(
				    tokenizer=__lowerCAmelCase  ,  mlm_probability=args.mlm_probability  ,  mlm=__lowerCAmelCase  ,  return_tensors='''tf'''      )

				def mask_with_collator(__lowerCAmelCase      ):
								# TF really needs an isin() function
								snake_case__	:      Optional[Any]     		=	(
								    ~tf.cast(batch['''attention_mask''']  ,  tf.bool      )
								    | (batch['''input_ids'''] == tokenizer.cls_token_id)
								    | (batch['''input_ids'''] == tokenizer.sep_token_id)
								)
								snake_case__		, snake_case__	:      int     		=	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

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

				snake_case__	:      str     		=	prepare_dataset(
				    __lowerCAmelCase  ,  decode_fn=__lowerCAmelCase  ,  mask_fn=__lowerCAmelCase  ,  batch_size=__lowerCAmelCase  ,  shuffle=__lowerCAmelCase  ,  shuffle_buffer_size=args.shuffle_buffer_size  ,  )

				snake_case__	:      int     		=	prepare_dataset(
				    __lowerCAmelCase  ,  decode_fn=__lowerCAmelCase  ,  mask_fn=__lowerCAmelCase  ,  batch_size=__lowerCAmelCase  ,  shuffle=__lowerCAmelCase  ,  )

				snake_case__	:      Union[str, Any]     		=	[]
				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__":
				A__								=						parse_args()
				main(args)






from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase      )   ->					Tuple:


				"""simple docstring"""

				snake_case__	:      Union[str, Any]     		=	[]
				for part_id in partition_order:
								snake_case__	:      Any     		=	df.where(f"""SPARK_PARTITION_ID() = {part_id}"""      ).collect()
								for row_idx, row in enumerate(__lowerCAmelCase      ):
												expected_row_ids_and_row_dicts.append((f"""{part_id}_{row_idx}""", row.asDict())      )
				return expected_row_ids_and_row_dicts



@require_not_windows
@require_dill_gt_0_3_2
def 			_lowerCAmelCase       (					)   ->					Tuple:


				"""simple docstring"""

				snake_case__	:      Any     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Optional[Any]     		=	spark.range(100      ).repartition(1      )
				snake_case__	:      Optional[int]     		=	Spark(__lowerCAmelCase      )
				# The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
				# that each partition can hold 2 rows.
				spark_builder._repartition_df_if_needed(max_shard_size=16      )
				# Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
				assert spark_builder.df.rdd.getNumPartitions() == 50



@require_not_windows
@require_dill_gt_0_3_2
def 			_lowerCAmelCase       (					)   ->					Optional[int]:


				"""simple docstring"""

				snake_case__	:      Any     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Dict     		=	spark.range(10      ).repartition(2      )
				snake_case__	:      Any     		=	[1, 0]
				snake_case__	:      Tuple     		=	_generate_iterable_examples(__lowerCAmelCase  ,  __lowerCAmelCase      )  # Reverse the partitions.
				snake_case__	:      Any     		=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase  ,  __lowerCAmelCase      )

				for i, (row_id, row_dict) in enumerate(generate_fn()      ):
								snake_case__		, snake_case__	:      Union[str, Any]     		=	expected_row_ids_and_row_dicts[i]
								assert row_id == expected_row_id
								assert row_dict == expected_row_dict



@require_not_windows
@require_dill_gt_0_3_2
def 			_lowerCAmelCase       (					)   ->					Any:


				"""simple docstring"""

				snake_case__	:      Any     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      List[Any]     		=	spark.range(10      ).repartition(1      )
				snake_case__	:      int     		=	SparkExamplesIterable(__lowerCAmelCase      )
				assert it.n_shards == 1
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase      ):
								assert row_id == f"""0_{i}"""
								assert row_dict == {"id": i}



@require_not_windows
@require_dill_gt_0_3_2
def 			_lowerCAmelCase       (					)   ->					Dict:


				"""simple docstring"""

				snake_case__	:      List[str]     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Tuple     		=	spark.range(30      ).repartition(3      )
				# Mock the generator so that shuffle reverses the partition indices.
				with patch('''numpy.random.Generator'''      ) as generator_mock:
								snake_case__	:      Union[str, Any]     		=	lambda __lowerCAmelCase      : x.reverse()
								snake_case__	:      List[str]     		=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase  ,  [2, 1, 0]      )

								snake_case__	:      List[str]     		=	SparkExamplesIterable(__lowerCAmelCase      ).shuffle_data_sources(__lowerCAmelCase      )
								assert shuffled_it.n_shards == 3
								for i, (row_id, row_dict) in enumerate(__lowerCAmelCase      ):
												snake_case__		, snake_case__	:      Optional[Any]     		=	expected_row_ids_and_row_dicts[i]
												assert row_id == expected_row_id
												assert row_dict == expected_row_dict



@require_not_windows
@require_dill_gt_0_3_2
def 			_lowerCAmelCase       (					)   ->					List[Any]:


				"""simple docstring"""

				snake_case__	:      Union[str, Any]     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Dict     		=	spark.range(20      ).repartition(4      )

				# Partitions 0 and 2
				snake_case__	:      List[Any]     		=	SparkExamplesIterable(__lowerCAmelCase      ).shard_data_sources(worker_id=0  ,  num_workers=2      )
				assert shard_it_a.n_shards == 2
				snake_case__	:      int     		=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase  ,  [0, 2]      )
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase      ):
								snake_case__		, snake_case__	:      Dict     		=	expected_row_ids_and_row_dicts_a[i]
								assert row_id == expected_row_id
								assert row_dict == expected_row_dict

				# Partitions 1 and 3
				snake_case__	:      List[str]     		=	SparkExamplesIterable(__lowerCAmelCase      ).shard_data_sources(worker_id=1  ,  num_workers=2      )
				assert shard_it_a.n_shards == 2
				snake_case__	:      List[Any]     		=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase  ,  [1, 3]      )
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase      ):
								snake_case__		, snake_case__	:      Optional[int]     		=	expected_row_ids_and_row_dicts_a[i]
								assert row_id == expected_row_id
								assert row_dict == expected_row_dict





@require_not_windows
@require_dill_gt_0_3_2
def 			_lowerCAmelCase       (					)   ->					Dict:


				"""simple docstring"""

				snake_case__	:      int     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Dict     		=	spark.range(100      ).repartition(1      )
				snake_case__	:      Tuple     		=	Spark(__lowerCAmelCase      )
				# Choose a small max_shard_size for maximum partitioning.
				spark_builder._repartition_df_if_needed(max_shard_size=1      )
				# The new number of partitions should not be greater than the number of rows.
				assert spark_builder.df.rdd.getNumPartitions() == 100

