
'''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 		_UpperCAmelCase   (	):






	"""simple docstring"""


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

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

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

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

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

	parser.add_argument(
	    "--max_length"			,							type=UpperCamelCase			,							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=UpperCamelCase			,							default=0.15			,							help="Fraction of tokens to mask during training."			,							)

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

	__lowerCAmelCase         =							parser.parse_args()
	return args






def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[Any]  ):






	"""simple docstring"""


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

	return tpu






def 		_UpperCAmelCase   (	UpperCamelCase:     List[str]  ):






	"""simple docstring"""


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

	return num_samples






def 		_UpperCAmelCase   (	UpperCamelCase:     List[str]			,							UpperCamelCase:     Union[str, Any]			,							UpperCamelCase:     Any			,							UpperCamelCase:     Tuple			,							UpperCamelCase:     List[Any]			,							UpperCamelCase:     str=None  ):






	"""simple docstring"""


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






def 		_UpperCAmelCase   (	UpperCamelCase:     Union[str, Any]  ):






	"""simple docstring"""


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

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

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

	__lowerCAmelCase         =							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         =							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         =							count_samples(UpperCamelCase  )

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

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

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

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

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

	__lowerCAmelCase         =							args.per_replica_batch_size * strategy.num_replicas_in_sync

	__lowerCAmelCase         =							prepare_dataset(
	    UpperCamelCase			,							decode_fn=UpperCamelCase			,							mask_fn=UpperCamelCase			,							batch_size=UpperCamelCase			,							shuffle=UpperCamelCase			,							shuffle_buffer_size=args.shuffle_buffer_size			,							)

	__lowerCAmelCase         =							prepare_dataset(
	    UpperCamelCase			,							decode_fn=UpperCamelCase			,							mask_fn=UpperCamelCase			,							batch_size=UpperCamelCase			,							shuffle=UpperCamelCase			,							)

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

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

	model.save_pretrained(args.output_dir  )


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








import os

import pytest

from datasets import (
    get_dataset_config_info,
    get_dataset_config_names,
    get_dataset_infos,
    get_dataset_split_names,
    inspect_dataset,
    inspect_metric,
)


UpperCamelCase_   	=	pytest.mark.integration






@pytest.mark.parametrize("path"			,							["paws", "csv"]  )
def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[int]			,							UpperCamelCase:     Tuple  ):






	"""simple docstring"""


	inspect_dataset(UpperCamelCase			,							UpperCamelCase  )
	__lowerCAmelCase         =							path + ".py"
	assert script_name in os.listdir(UpperCamelCase  )
	assert "__pycache__" not in os.listdir(UpperCamelCase  )






@pytest.mark.filterwarnings("ignore:inspect_metric is deprecated:FutureWarning"  )
@pytest.mark.filterwarnings("ignore:metric_module_factory is deprecated:FutureWarning"  )
@pytest.mark.parametrize("path"			,							["accuracy"]  )
def 		_UpperCAmelCase   (	UpperCamelCase:     List[Any]			,							UpperCamelCase:     Any  ):






	"""simple docstring"""


	inspect_metric(UpperCamelCase			,							UpperCamelCase  )
	__lowerCAmelCase         =							path + ".py"
	assert script_name in os.listdir(UpperCamelCase  )
	assert "__pycache__" not in os.listdir(UpperCamelCase  )






@pytest.mark.parametrize(
    "path, config_name, expected_splits"			,							[
        ("squad", "plain_text", ["train", "validation"]),
        ("dalle-mini/wit", "dalle-mini--wit", ["train"]),
        ("paws", "labeled_final", ["train", "test", "validation"]),
    ]			,							)
def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[Any]			,							UpperCamelCase:     int			,							UpperCamelCase:     Tuple  ):






	"""simple docstring"""


	__lowerCAmelCase         =							get_dataset_config_info(UpperCamelCase			,							config_name=UpperCamelCase  )
	assert info.config_name == config_name
	assert list(info.splits.keys()  ) == expected_splits






@pytest.mark.parametrize(
    "path, config_name, expected_exception"			,							[
        ("paws", None, ValueError),
    ]			,							)
def 		_UpperCAmelCase   (	UpperCamelCase:     Dict			,							UpperCamelCase:     Any			,							UpperCamelCase:     List[str]  ):






	"""simple docstring"""


	with pytest.raises(UpperCamelCase  ):
		get_dataset_config_info(UpperCamelCase			,							config_name=UpperCamelCase  )






@pytest.mark.parametrize(
    "path, expected"			,							[
        ("squad", "plain_text"),
        ("acronym_identification", "default"),
        ("lhoestq/squad", "plain_text"),
        ("lhoestq/test", "default"),
        ("lhoestq/demo1", "lhoestq--demo1"),
        ("dalle-mini/wit", "dalle-mini--wit"),
    ]			,							)
def 		_UpperCAmelCase   (	UpperCamelCase:     List[str]			,							UpperCamelCase:     Optional[int]  ):






	"""simple docstring"""


	__lowerCAmelCase         =							get_dataset_config_names(UpperCamelCase  )
	assert expected in config_names






@pytest.mark.parametrize(
    "path, expected_configs, expected_splits_in_first_config"			,							[
        ("squad", ["plain_text"], ["train", "validation"]),
        ("dalle-mini/wit", ["dalle-mini--wit"], ["train"]),
        ("paws", ["labeled_final", "labeled_swap", "unlabeled_final"], ["train", "test", "validation"]),
    ]			,							)
def 		_UpperCAmelCase   (	UpperCamelCase:     str			,							UpperCamelCase:     Tuple			,							UpperCamelCase:     List[Any]  ):






	"""simple docstring"""


	__lowerCAmelCase         =							get_dataset_infos(UpperCamelCase  )
	assert list(infos.keys()  ) == expected_configs
	__lowerCAmelCase         =							expected_configs[0]
	assert expected_config in infos
	__lowerCAmelCase         =							infos[expected_config]
	assert info.config_name == expected_config
	assert list(info.splits.keys()  ) == expected_splits_in_first_config






@pytest.mark.parametrize(
    "path, expected_config, expected_splits"			,							[
        ("squad", "plain_text", ["train", "validation"]),
        ("dalle-mini/wit", "dalle-mini--wit", ["train"]),
        ("paws", "labeled_final", ["train", "test", "validation"]),
    ]			,							)
def 		_UpperCAmelCase   (	UpperCamelCase:     Optional[int]			,							UpperCamelCase:     int			,							UpperCamelCase:     Optional[Any]  ):






	"""simple docstring"""


	__lowerCAmelCase         =							get_dataset_infos(UpperCamelCase  )
	assert expected_config in infos
	__lowerCAmelCase         =							infos[expected_config]
	assert info.config_name == expected_config
	assert list(info.splits.keys()  ) == expected_splits







@pytest.mark.parametrize(
    "path, config_name, expected_exception"			,							[
        ("paws", None, ValueError),
    ]			,							)
def 		_UpperCAmelCase   (	UpperCamelCase:     Dict			,							UpperCamelCase:     List[str]			,							UpperCamelCase:     List[Any]  ):






	"""simple docstring"""


	with pytest.raises(UpperCamelCase  ):
		get_dataset_split_names(UpperCamelCase			,							config_name=UpperCamelCase  )






