






'''simple docstring'''



import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
 import tensorflow as tf
 from tensorflow.python.framework.errors_impl import ResourceExhaustedError

 from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
 import pyanvml.pyanvml as nvml

__UpperCamelCase   :							str         =   logging.get_logger(__name__)



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						bool,       SCREAMING_SNAKE_CASE__:						bool     )	->      List[Any]:

   """simple docstring"""
   def run_func(SCREAMING_SNAKE_CASE__:						Optional[int]     ):
      @wraps(SCREAMING_SNAKE_CASE__     )
      def run_in_eager_mode(*SCREAMING_SNAKE_CASE__:						Tuple,       **SCREAMING_SNAKE_CASE__:						Optional[int]     ):
         return func(*SCREAMING_SNAKE_CASE__,       **SCREAMING_SNAKE_CASE__     )

      @wraps(SCREAMING_SNAKE_CASE__     )
      @tf.function(experimental_compile=SCREAMING_SNAKE_CASE__     )
      def run_in_graph_mode(*SCREAMING_SNAKE_CASE__:						int,       **SCREAMING_SNAKE_CASE__:						Dict     ):
         return func(*SCREAMING_SNAKE_CASE__,       **SCREAMING_SNAKE_CASE__     )

      if do_eager_mode is True:
         if use_xla is not False:
            raise ValueError(
                'Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'     )
         return run_in_eager_mode
      else:
         return run_in_graph_mode

   return run_func



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						int,       SCREAMING_SNAKE_CASE__:						int,       SCREAMING_SNAKE_CASE__:						int     )	->      ["tf.Tensor"]:

   """simple docstring"""
   __a					  =     random.Random()
   __a					  =     [rng.randint(0,       vocab_size - 1     ) for i in range(batch_size * sequence_length     )]
   return tf.constant(SCREAMING_SNAKE_CASE__,       shape=(batch_size, sequence_length),       dtype=tf.intaa     )



class 				__SCREAMING_SNAKE_CASE							(				_lowerCAmelCase  ):
      __a	       =42
      __a	       =42
      __a	       ="TensorFlow"




      @property
      def 			__UpperCamelCase				(	self   )				->int:



         '''simple docstring'''
         return tf.__version__




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->float:



         '''simple docstring'''
         # initialize GPU on separate process
         __a					  =     self.args.strategy
         if strategy is None:
            raise ValueError('A device strategy has to be initialized before using TensorFlow.'   )
         __a					  =     self._prepare_inference_func(lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )
         return self._measure_speed(_inference   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->float:



         '''simple docstring'''
         __a					  =     self.args.strategy
         if strategy is None:
            raise ValueError('A device strategy has to be initialized before using TensorFlow.'   )
         __a					  =     self._prepare_train_func(lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )
         return self._measure_speed(_train   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->[Memory, Optional[MemorySummary]]:



         '''simple docstring'''
         # initialize GPU on separate process
         if self.args.is_gpu:
            tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]			, lowerCamelCase   )
         __a					  =     self.args.strategy
         if strategy is None:
            raise ValueError('A device strategy has to be initialized before using TensorFlow.'   )
         __a					  =     self._prepare_inference_func(lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )
         return self._measure_memory(_inference   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->[Memory, Optional[MemorySummary]]:



         '''simple docstring'''
         if self.args.is_gpu:
            tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]			, lowerCamelCase   )
         __a					  =     self.args.strategy
         if strategy is None:
            raise ValueError('A device strategy has to be initialized before using TensorFlow.'   )

         __a					  =     self._prepare_train_func(lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )
         return self._measure_memory(_train   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->Callable[[], None]:



         '''simple docstring'''
         __a					  =     self.config_dict[model_name]

         if self.args.fpaa:
            raise NotImplementedError('Mixed precision is currently not supported.'   )

         __a					  =     (
             hasattr(lowerCamelCase			, 'architectures'   )
             and isinstance(config.architectures			, lowerCamelCase   )
             and len(config.architectures   ) > 0
         )
         if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
               __a					  =     'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
               __a					  =     __import__('transformers'			, fromlist=[model_class]   )
               __a					  =     getattr(lowerCamelCase			, lowerCamelCase   )
               __a					  =     model_cls(lowerCamelCase   )
            except ImportError:
               raise ImportError(
                   F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                   ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'   )
         else:
            __a					  =     TF_MODEL_MAPPING[config.__class__](lowerCamelCase   )

         # encoder-decoder has vocab size saved differently
         __a					  =     config.vocab_size if hasattr(lowerCamelCase			, 'vocab_size'   ) else config.encoder.vocab_size
         __a					  =     random_input_ids(lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )

         @run_with_tf_optimizations(self.args.eager_mode			, self.args.use_xla   )
         def encoder_decoder_forward():
            return model(lowerCamelCase			, decoder_input_ids=lowerCamelCase			, training=lowerCamelCase   )

         @run_with_tf_optimizations(self.args.eager_mode			, self.args.use_xla   )
         def encoder_forward():
            return model(lowerCamelCase			, training=lowerCamelCase   )

         __a					  =     encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

         return _inference




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->Callable[[], None]:



         '''simple docstring'''
         __a					  =     self.config_dict[model_name]

         if self.args.eager_mode is not False:
            raise ValueError('Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.'   )

         if self.args.fpaa:
            raise NotImplementedError('Mixed precision is currently not supported.'   )

         __a					  =     (
             hasattr(lowerCamelCase			, 'architectures'   )
             and isinstance(config.architectures			, lowerCamelCase   )
             and len(config.architectures   ) > 0
         )
         if not self.args.only_pretrain_model and has_model_class_in_config:
            try:
               __a					  =     'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
               __a					  =     __import__('transformers'			, fromlist=[model_class]   )
               __a					  =     getattr(lowerCamelCase			, lowerCamelCase   )
               __a					  =     model_cls(lowerCamelCase   )
            except ImportError:
               raise ImportError(
                   F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                   ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'   )
         else:
            __a					  =     TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowerCamelCase   )

         # encoder-decoder has vocab size saved differently
         __a					  =     config.vocab_size if hasattr(lowerCamelCase			, 'vocab_size'   ) else config.encoder.vocab_size
         __a					  =     random_input_ids(lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )

         @run_with_tf_optimizations(self.args.eager_mode			, self.args.use_xla   )
         def encoder_decoder_train():
            __a					  =     model(lowerCamelCase			, decoder_input_ids=lowerCamelCase			, labels=lowerCamelCase			, training=lowerCamelCase   )[0]
            __a					  =     tf.gradients(lowerCamelCase			, model.trainable_variables   )
            return gradients

         @run_with_tf_optimizations(self.args.eager_mode			, self.args.use_xla   )
         def encoder_train():
            __a					  =     model(lowerCamelCase			, labels=lowerCamelCase			, training=lowerCamelCase   )[0]
            __a					  =     tf.gradients(lowerCamelCase			, model.trainable_variables   )
            return gradients

         __a					  =     encoder_decoder_train if config.is_encoder_decoder else encoder_train

         return _train




      def 			__UpperCamelCase				(	self			, lowerCamelCase   )				->float:



         '''simple docstring'''
         with self.args.strategy.scope():
            try:
               if self.args.is_tpu or self.args.use_xla:
                  # run additional 10 times to stabilize compilation for tpu
                  logger.info('Do inference on TPU. Running model 5 times to stabilize compilation'   )
                  timeit.repeat(lowerCamelCase			, repeat=1			, number=5   )

               # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
               __a					  =     timeit.repeat(
                   lowerCamelCase			, repeat=self.args.repeat			, number=10			, )

               return min(lowerCamelCase   ) / 10.0
            except ResourceExhaustedError as e:
               self.print_fn(F"""Doesn't fit on GPU. {e}"""   )






      def 			__UpperCamelCase				(	self			, lowerCamelCase   )				->[Memory, MemorySummary]:



         '''simple docstring'''
         logger.info(
             'Note that TensorFlow allocates more memory than '
             'it might need to speed up computation. '
             'The memory reported here corresponds to the memory '
             'reported by `nvidia-smi`, which can vary depending '
             'on total available memory on the GPU that is used.'   )
         with self.args.strategy.scope():
            try:
               if self.args.trace_memory_line_by_line:
                  if not self.args.eager_mode:
                     raise ValueError(
                         '`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'
                         ' consumption line by line.'   )
                  __a					  =     start_memory_tracing('transformers'   )

               if self.args.is_tpu:
                  # tpu
                  raise NotImplementedError(
                      'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'
                      ' with `args.memory=False`'   )
               elif self.args.is_gpu:
                  # gpu
                  if not is_pyanvml_available():
                     logger.warning(
                         'py3nvml not installed, we won\'t log GPU memory usage. '
                         'Install py3nvml (pip install py3nvml) to log information about GPU.'   )
                     __a					  =     'N/A'
                  else:
                     logger.info(
                         'Measuring total GPU usage on GPU device. Make sure to not have additional processes'
                         ' running on the same GPU.'   )
                     # init nvml
                     nvml.nvmlInit()
                     func()
                     __a					  =     nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx   )
                     __a					  =     nvml.nvmlDeviceGetMemoryInfo(lowerCamelCase   )
                     __a					  =     meminfo.used
                     __a					  =     Memory(lowerCamelCase   )
                     # shutdown nvml
                     nvml.nvmlShutdown()
               else:
                  # cpu
                  if self.args.trace_memory_line_by_line:
                     logger.info(
                         'When enabling line by line tracing, the max peak memory for CPU is inaccurate in'
                         ' TensorFlow.'   )
                     __a					  =     None
                  else:
                     __a					  =     measure_peak_memory_cpu(lowerCamelCase   )
                     __a					  =     Memory(lowerCamelCase   ) if isinstance(lowerCamelCase			, lowerCamelCase   ) else memory_bytes
               if self.args.trace_memory_line_by_line:
                  __a					  =     stop_memory_tracing(lowerCamelCase   )
                  if memory is None:
                     __a					  =     summary.total
               else:
                  __a					  =     None

               return memory, summary
            except ResourceExhaustedError as e:
               self.print_fn(F"""Doesn't fit on GPU. {e}"""   )
               return "N/A", None







'''simple docstring'''



from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import fa_score, matthews_corrcoef

import datasets


__UpperCamelCase   :							Optional[int]         =   """\
@inproceedings{wang2019glue,
  title={{GLUE}: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding},
  author={Wang, Alex and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R.},
  note={In the Proceedings of ICLR.},
  year={2019}
}
"""

__UpperCamelCase   :							Tuple         =   """\
GLUE, the General Language Understanding Evaluation benchmark
(https://gluebenchmark.com/) is a collection of resources for training,
evaluating, and analyzing natural language understanding systems.
"""

__UpperCamelCase   :							Optional[int]         =   """
Compute GLUE evaluation metric associated to each GLUE dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
Returns: depending on the GLUE subset, one or several of:
    \"accuracy\": Accuracy
    \"f1\": F1 score
    \"pearson\": Pearson Correlation
    \"spearmanr\": Spearman Correlation
    \"matthews_correlation\": Matthew Correlation
Examples:

    >>> glue_metric = datasets.load_metric('glue', 'sst2')  # 'sst2' or any of [\"mnli\", \"mnli_mismatched\", \"mnli_matched\", \"qnli\", \"rte\", \"wnli\", \"hans\"]
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'accuracy': 1.0}

    >>> glue_metric = datasets.load_metric('glue', 'mrpc')  # 'mrpc' or 'qqp'
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'accuracy': 1.0, 'f1': 1.0}

    >>> glue_metric = datasets.load_metric('glue', 'stsb')
    >>> references = [0., 1., 2., 3., 4., 5.]
    >>> predictions = [0., 1., 2., 3., 4., 5.]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print({\"pearson\": round(results[\"pearson\"], 2), \"spearmanr\": round(results[\"spearmanr\"], 2)})
    {'pearson': 1.0, 'spearmanr': 1.0}

    >>> glue_metric = datasets.load_metric('glue', 'cola')
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'matthews_correlation': 1.0}
"""



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						Optional[int],       SCREAMING_SNAKE_CASE__:						str     )	->      List[Any]:

   """simple docstring"""
   return float((preds == labels).mean()     )



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						Optional[Any],       SCREAMING_SNAKE_CASE__:						List[str]     )	->      Optional[Any]:

   """simple docstring"""
   __a					  =     simple_accuracy(SCREAMING_SNAKE_CASE__,       SCREAMING_SNAKE_CASE__     )
   __a					  =     float(fa_score(y_true=SCREAMING_SNAKE_CASE__,       y_pred=SCREAMING_SNAKE_CASE__     )     )
   return {
       "accuracy": acc,
       "f1": fa,
   }



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						int,       SCREAMING_SNAKE_CASE__:						Optional[Any]     )	->      int:

   """simple docstring"""
   __a					  =     float(pearsonr(SCREAMING_SNAKE_CASE__,       SCREAMING_SNAKE_CASE__     )[0]     )
   __a					  =     float(spearmanr(SCREAMING_SNAKE_CASE__,       SCREAMING_SNAKE_CASE__     )[0]     )
   return {
       "pearson": pearson_corr,
       "spearmanr": spearman_corr,
   }



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




      def 			__UpperCamelCase				(	self   )				->int:



         '''simple docstring'''
         if self.config_name not in [
             "sst2",
             "mnli",
             "mnli_mismatched",
             "mnli_matched",
             "cola",
             "stsb",
             "mrpc",
             "qqp",
             "qnli",
             "rte",
             "wnli",
             "hans",
         ]:
            raise KeyError(
                'You should supply a configuration name selected in '
                '["sst2", "mnli", "mnli_mismatched", "mnli_matched", '
                '"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]'   )
         return datasets.MetricInfo(
             description=_DESCRIPTION			, citation=_CITATION			, inputs_description=_KWARGS_DESCRIPTION			, features=datasets.Features(
                 {
                     'predictions': datasets.Value('int64' if self.config_name != 'stsb' else 'float32'   ),
                     'references': datasets.Value('int64' if self.config_name != 'stsb' else 'float32'   ),
                 }   )			, codebase_urls=[]			, reference_urls=[]			, format='numpy'			, )






      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase   )				->Tuple:



         '''simple docstring'''
         if self.config_name == "cola":
            return {"matthews_correlation": matthews_corrcoef(lowerCamelCase			, lowerCamelCase   )}
         elif self.config_name == "stsb":
            return pearson_and_spearman(lowerCamelCase			, lowerCamelCase   )
         elif self.config_name in ["mrpc", "qqp"]:
            return acc_and_fa(lowerCamelCase			, lowerCamelCase   )
         elif self.config_name in ["sst2", "mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]:
            return {"accuracy": simple_accuracy(lowerCamelCase			, lowerCamelCase   )}
         else:
            raise KeyError(
                'You should supply a configuration name selected in '
                '["sst2", "mnli", "mnli_mismatched", "mnli_matched", '
                '"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]'   )