



"""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       :					Union[str, Any]      					=					logging.getLogger(__name__)

__lowerCAmelCase       :					Optional[int]      					=					tf.data.AUTOTUNE
def 	__lowerCAmelCase						(		):


						'''simple docstring'''






						snake_case_   : List[str]						=			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."""    )

						snake_case_   : List[Any]						=			parser.parse_args()
						return args
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Any    ):


						'''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_   : Dict						=			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 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[int]    ):


						'''simple docstring'''






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

						return num_samples
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[Any]					, __UpperCamelCase      :			Tuple					, __UpperCamelCase      :			List[str]					, __UpperCamelCase      :			int					, __UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			Any=None    ):


						'''simple docstring'''






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


						'''simple docstring'''






						if not args.no_tpu:
												snake_case_   : List[Any]						=			initialize_tpu(__UpperCamelCase    )
												snake_case_   : Dict						=			tf.distribute.TPUStrategy(__UpperCamelCase    )
						else:
												snake_case_   : Union[str, Any]						=			tf.distribute.OneDeviceStrategy(device="""/gpu:0"""    )

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

						snake_case_   : Union[str, Any]						=			AutoTokenizer.from_pretrained(args.tokenizer    )
						snake_case_   : Optional[int]						=			AutoConfig.from_pretrained(args.pretrained_model_config    )
						snake_case_   : int						=			tokenizer.vocab_size

						snake_case_   : List[str]						=			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_   : Tuple						=			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_   : Optional[int]						=			count_samples(__UpperCamelCase    )

						snake_case_   : Optional[int]						=			num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
						snake_case_   : Dict						=			steps_per_epoch * args.num_epochs

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

						def mask_with_collator(__UpperCamelCase      :			Tuple    ):
												# 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_   : int						=			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

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

						snake_case_   : Any						=			prepare_dataset(
						    __UpperCamelCase					, decode_fn=__UpperCamelCase					, mask_fn=__UpperCamelCase					, batch_size=__UpperCamelCase					, shuffle=__UpperCamelCase					, shuffle_buffer_size=args.shuffle_buffer_size					, )

						snake_case_   : List[Any]						=			prepare_dataset(
						    __UpperCamelCase					, decode_fn=__UpperCamelCase					, mask_fn=__UpperCamelCase					, batch_size=__UpperCamelCase					, shuffle=__UpperCamelCase					, )

						snake_case_   : List[str]						=			[]
						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__":
					__lowerCAmelCase       :					int      					=					parse_args()
					main(args)








"""simple docstring"""





from sklearn.metrics import fa_score, matthews_corrcoef

import datasets

from .record_evaluation import evaluate as evaluate_record


__lowerCAmelCase       :					List[str]      					=					'''\
@article{wang2019superglue,
  title={SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems},
  author={Wang, Alex and Pruksachatkun, Yada and Nangia, Nikita and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R},
  journal={arXiv preprint arXiv:1905.00537},
  year={2019}
}
'''

__lowerCAmelCase       :					Optional[Any]      					=					'''\
SuperGLUE (https://super.gluebenchmark.com/) is a new benchmark styled after
GLUE with a new set of more difficult language understanding tasks, improved
resources, and a new public leaderboard.
'''

__lowerCAmelCase       :					str      					=					'''
Compute SuperGLUE evaluation metric associated to each SuperGLUE dataset.
Args:
    predictions: list of predictions to score. Depending on the SuperGlUE subset:
        - for \'record\': list of question-answer dictionaries with the following keys:
            - \'idx\': index of the question as specified by the dataset
            - \'prediction_text\': the predicted answer text
        - for \'multirc\': list of question-answer dictionaries with the following keys:
            - \'idx\': index of the question-answer pair as specified by the dataset
            - \'prediction\': the predicted answer label
        - otherwise: list of predicted labels
    references: list of reference labels. Depending on the SuperGLUE subset:
        - for \'record\': list of question-answers dictionaries with the following keys:
            - \'idx\': index of the question as specified by the dataset
            - \'answers\': list of possible answers
        - otherwise: list of reference labels
Returns: depending on the SuperGLUE subset:
    - for \'record\':
        - \'exact_match\': Exact match between answer and gold answer
        - \'f1\': F1 score
    - for \'multirc\':
        - \'exact_match\': Exact match between answer and gold answer
        - \'f1_m\': Per-question macro-F1 score
        - \'f1_a\': Average F1 score over all answers
    - for \'axb\':
        \'matthews_correlation\': Matthew Correlation
    - for \'cb\':
        - \'accuracy\': Accuracy
        - \'f1\': F1 score
    - for all others:
        - \'accuracy\': Accuracy
Examples:

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'copa\')  # any of ["copa", "rte", "wic", "wsc", "wsc.fixed", "boolq", "axg"]
    >>> predictions = [0, 1]
    >>> references = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'cb\')
    >>> predictions = [0, 1]
    >>> references = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0, \'f1\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'record\')
    >>> predictions = [{\'idx\': {\'passage\': 0, \'query\': 0}, \'prediction_text\': \'answer\'}]
    >>> references = [{\'idx\': {\'passage\': 0, \'query\': 0}, \'answers\': [\'answer\', \'another_answer\']}]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact_match\': 1.0, \'f1\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'multirc\')
    >>> predictions = [{\'idx\': {\'answer\': 0, \'paragraph\': 0, \'question\': 0}, \'prediction\': 0}, {\'idx\': {\'answer\': 1, \'paragraph\': 2, \'question\': 3}, \'prediction\': 1}]
    >>> references = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact_match\': 1.0, \'f1_m\': 1.0, \'f1_a\': 1.0}

    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'axb\')
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = super_glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'matthews_correlation\': 1.0}
'''
def 	__lowerCAmelCase						(		__UpperCamelCase      :			List[str]					, __UpperCamelCase      :			Any    ):


						'''simple docstring'''






						return float((preds == labels).mean()    )
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Any					, __UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			str="binary"    ):


						'''simple docstring'''






						snake_case_   : Optional[Any]						=			simple_accuracy(__UpperCamelCase					, __UpperCamelCase    )
						snake_case_   : Dict						=			float(fa_score(y_true=__UpperCamelCase					, y_pred=__UpperCamelCase					, average=__UpperCamelCase    )    )
						return {
						    "accuracy": acc,
						    "f1": fa,
						}
def 	__lowerCAmelCase						(		__UpperCamelCase      :			str					, __UpperCamelCase      :			int    ):


						'''simple docstring'''






						snake_case_   : List[Any]						=			{}
						for id_pred, label in zip(__UpperCamelCase					, __UpperCamelCase    ):
												snake_case_   : Optional[int]						=			F'{id_pred["idx"]["paragraph"]}-{id_pred["idx"]["question"]}'
												snake_case_   : Union[str, Any]						=			id_pred["""prediction"""]
												if question_id in question_map:
																		question_map[question_id].append((pred, label)    )
												else:
																		snake_case_   : str						=			[(pred, label)]
						snake_case_			,		snake_case_   : List[str]						=			[], []
						for question, preds_labels in question_map.items():
												snake_case_			,		snake_case_   : Optional[Any]						=			zip(*__UpperCamelCase    )
												snake_case_   : int						=			fa_score(y_true=__UpperCamelCase					, y_pred=__UpperCamelCase					, average="""macro"""    )
												fas.append(__UpperCamelCase    )
												snake_case_   : Dict						=			int(sum(pred == label for pred, label in preds_labels    ) == len(__UpperCamelCase    )    )
												ems.append(__UpperCamelCase    )
						snake_case_   : Optional[int]						=			float(sum(__UpperCamelCase    ) / len(__UpperCamelCase    )    )
						snake_case_   : Any						=			sum(__UpperCamelCase    ) / len(__UpperCamelCase    )
						snake_case_   : int						=			float(fa_score(y_true=__UpperCamelCase					, y_pred=[id_pred["""prediction"""] for id_pred in ids_preds]    )    )
						return {"exact_match": em, "f1_m": fa_m, "f1_a": fa_a}






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,				_KWARGS_DESCRIPTION		)
class    _lowerCAmelCase    (				datasets.Metric		):





		"""simple docstring"""






		def    UpperCAmelCase__							(					self   ) ->						Optional[int]:







								'''simple docstring'''

								if self.config_name not in [
								    "boolq",
								    "cb",
								    "copa",
								    "multirc",
								    "record",
								    "rte",
								    "wic",
								    "wsc",
								    "wsc.fixed",
								    "axb",
								    "axg",
								]:
														raise KeyError(
														    """You should supply a configuration name selected in """
														    """[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]"""   )
								return datasets.MetricInfo(
								    description=_DESCRIPTION							,    citation=_CITATION							,    inputs_description=_KWARGS_DESCRIPTION							,    features=datasets.Features(self._get_feature_types()   )							,    codebase_urls=[]							,    reference_urls=[]							,    format="""numpy""" if not self.config_name == """record""" and not self.config_name == """multirc""" else None							,    )





		def    UpperCAmelCase__							(					self   ) ->						Optional[int]:







								'''simple docstring'''

								if self.config_name == "record":
														return {
														    "predictions": {
														        "idx": {
														            "passage": datasets.Value("""int64"""   ),
														            "query": datasets.Value("""int64"""   ),
														        },
														        "prediction_text": datasets.Value("""string"""   ),
														    },
														    "references": {
														        "idx": {
														            "passage": datasets.Value("""int64"""   ),
														            "query": datasets.Value("""int64"""   ),
														        },
														        "answers": datasets.Sequence(datasets.Value("""string"""   )   ),
														    },
														}
								elif self.config_name == "multirc":
														return {
														    "predictions": {
														        "idx": {
														            "answer": datasets.Value("""int64"""   ),
														            "paragraph": datasets.Value("""int64"""   ),
														            "question": datasets.Value("""int64"""   ),
														        },
														        "prediction": datasets.Value("""int64"""   ),
														    },
														    "references": datasets.Value("""int64"""   ),
														}
								else:
														return {
														    "predictions": datasets.Value("""int64"""   ),
														    "references": datasets.Value("""int64"""   ),
														}





		def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase   ) ->						List[str]:







								'''simple docstring'''

								if self.config_name == "axb":
														return {"matthews_correlation": matthews_corrcoef(_lowercase							,    _lowercase   )}
								elif self.config_name == "cb":
														return acc_and_fa(_lowercase							,    _lowercase							,    fa_avg="""macro"""   )
								elif self.config_name == "record":
														snake_case_   : Optional[Any]						=			[
														    {
														        """qas""": [
														            {"""id""": ref["""idx"""]["""query"""], """answers""": [{"""text""": ans} for ans in ref["""answers"""]]}
														            for ref in references
														        ]
														    }
														]
														snake_case_   : Dict						=			{pred["""idx"""]["""query"""]: pred["""prediction_text"""] for pred in predictions}
														return evaluate_record(_lowercase							,    _lowercase   )[0]
								elif self.config_name == "multirc":
														return evaluate_multirc(_lowercase							,    _lowercase   )
								elif self.config_name in ["copa", "rte", "wic", "wsc", "wsc.fixed", "boolq", "axg"]:
														return {"accuracy": simple_accuracy(_lowercase							,    _lowercase   )}
								else:
														raise KeyError(
														    """You should supply a configuration name selected in """
														    """[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]"""   )

