






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

A        =			logging.get_logger(__name__)



def 					_lowerCamelCase(			lowerCAmelCase__				:		bool    ,    lowerCAmelCase__				:		bool      ):





				'''simple docstring'''







				def run_func(lowerCAmelCase__				:		int      ):
								@wraps(lowerCAmelCase__      )
								def run_in_eager_mode(*lowerCAmelCase__				:		Tuple    ,    **lowerCAmelCase__				:		str      ):
												return func(*lowerCAmelCase__    ,    **lowerCAmelCase__      )

								@wraps(lowerCAmelCase__      )
								@tf.function(experimental_compile=lowerCAmelCase__      )
								def run_in_graph_mode(*lowerCAmelCase__				:		Union[str, Any]    ,    **lowerCAmelCase__				:		Union[str, Any]      ):
												return func(*lowerCAmelCase__    ,    **lowerCAmelCase__      )

								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 					_lowerCamelCase(			lowerCAmelCase__				:		int    ,    lowerCAmelCase__				:		int    ,    lowerCAmelCase__				:		int      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   random.Random()
				SCREAMING_SNAKE_CASE_   :	Tuple													=   [rng.randint(0    ,    vocab_size - 1      ) for i in range(batch_size * sequence_length      )]
				return tf.constant(lowerCAmelCase__    ,    shape=(batch_size, sequence_length)    ,    dtype=tf.intaa      )


class 							__a							(  __A						):





		'''simple docstring'''



		UpperCAmelCase__      :				TensorFlowBenchmarkArguments
		UpperCAmelCase__      :				PretrainedConfig
		UpperCAmelCase__      :				str				    =						"TensorFlow"


		@property
		def   __snake_case			(						self       ):
						return tf.__version__


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						# initialize GPU on separate process
						SCREAMING_SNAKE_CASE_   :	List[str]													=   self.args.strategy
						if strategy is None:
										raise ValueError('A device strategy has to be initialized before using TensorFlow.'       )
						SCREAMING_SNAKE_CASE_   :	Tuple													=   self._prepare_inference_func(UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       )
						return self._measure_speed(_inference       )


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   self.args.strategy
						if strategy is None:
										raise ValueError('A device strategy has to be initialized before using TensorFlow.'       )
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   self._prepare_train_func(UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       )
						return self._measure_speed(_train       )


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						# initialize GPU on separate process
						if self.args.is_gpu:
										tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx] ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   self.args.strategy
						if strategy is None:
										raise ValueError('A device strategy has to be initialized before using TensorFlow.'       )
						SCREAMING_SNAKE_CASE_   :	Tuple													=   self._prepare_inference_func(UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       )
						return self._measure_memory(_inference       )


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						if self.args.is_gpu:
										tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx] ,    UpperCamelCase__       )
						SCREAMING_SNAKE_CASE_   :	List[str]													=   self.args.strategy
						if strategy is None:
										raise ValueError('A device strategy has to be initialized before using TensorFlow.'       )

						SCREAMING_SNAKE_CASE_   :	List[Any]													=   self._prepare_train_func(UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       )
						return self._measure_memory(_train       )


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   self.config_dict[model_name]

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

						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   (
						    hasattr(UpperCamelCase__ ,    'architectures'       )
						    and isinstance(config.architectures ,    UpperCamelCase__       )
						    and len(config.architectures       ) > 0
						)
						if not self.args.only_pretrain_model and has_model_class_in_config:
										try:
														SCREAMING_SNAKE_CASE_   :	List[str]													=   'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
														SCREAMING_SNAKE_CASE_   :	Tuple													=   __import__('transformers' ,    fromlist=[model_class]       )
														SCREAMING_SNAKE_CASE_   :	List[str]													=   getattr(UpperCamelCase__ ,    UpperCamelCase__       )
														SCREAMING_SNAKE_CASE_   :	Tuple													=   model_cls(UpperCamelCase__       )
										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:
										SCREAMING_SNAKE_CASE_   :	int													=   TF_MODEL_MAPPING[config.__class__](UpperCamelCase__       )

						# encoder-decoder has vocab size saved differently
						SCREAMING_SNAKE_CASE_   :	Any													=   config.vocab_size if hasattr(UpperCamelCase__ ,    'vocab_size'       ) else config.encoder.vocab_size
						SCREAMING_SNAKE_CASE_   :	str													=   random_input_ids(UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       )

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

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

						SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

						return _inference


		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						SCREAMING_SNAKE_CASE_   :	List[str]													=   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.'       )

						SCREAMING_SNAKE_CASE_   :	Tuple													=   (
						    hasattr(UpperCamelCase__ ,    'architectures'       )
						    and isinstance(config.architectures ,    UpperCamelCase__       )
						    and len(config.architectures       ) > 0
						)
						if not self.args.only_pretrain_model and has_model_class_in_config:
										try:
														SCREAMING_SNAKE_CASE_   :	Dict													=   'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
														SCREAMING_SNAKE_CASE_   :	str													=   __import__('transformers' ,    fromlist=[model_class]       )
														SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   getattr(UpperCamelCase__ ,    UpperCamelCase__       )
														SCREAMING_SNAKE_CASE_   :	Dict													=   model_cls(UpperCamelCase__       )
										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:
										SCREAMING_SNAKE_CASE_   :	Optional[int]													=   TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](UpperCamelCase__       )

						# encoder-decoder has vocab size saved differently
						SCREAMING_SNAKE_CASE_   :	int													=   config.vocab_size if hasattr(UpperCamelCase__ ,    'vocab_size'       ) else config.encoder.vocab_size
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   random_input_ids(UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       )

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

						@run_with_tf_optimizations(self.args.eager_mode ,    self.args.use_xla       )
						def encoder_train():
										SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   model(UpperCamelCase__ ,    labels=UpperCamelCase__ ,    training=UpperCamelCase__       )[0]
										SCREAMING_SNAKE_CASE_   :	Optional[int]													=   tf.gradients(UpperCamelCase__ ,    model.trainable_variables       )
										return gradients

						SCREAMING_SNAKE_CASE_   :	Tuple													=   encoder_decoder_train if config.is_encoder_decoder else encoder_train

						return _train


		def   __snake_case			(						self ,    UpperCamelCase__       ):
						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(UpperCamelCase__ ,    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
														SCREAMING_SNAKE_CASE_   :	Optional[int]													=   timeit.repeat(
														    UpperCamelCase__ ,    repeat=self.args.repeat ,    number=10 ,    )

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






		def   __snake_case			(						self ,    UpperCamelCase__       ):
						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.'       )
																		SCREAMING_SNAKE_CASE_   :	Tuple													=   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.'       )
																						SCREAMING_SNAKE_CASE_   :	int													=   '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()
																						SCREAMING_SNAKE_CASE_   :	Tuple													=   nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx       )
																						SCREAMING_SNAKE_CASE_   :	int													=   nvml.nvmlDeviceGetMemoryInfo(UpperCamelCase__       )
																						SCREAMING_SNAKE_CASE_   :	Optional[int]													=   meminfo.used
																						SCREAMING_SNAKE_CASE_   :	Tuple													=   Memory(UpperCamelCase__       )
																						# 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.'       )
																						SCREAMING_SNAKE_CASE_   :	Any													=   None
																		else:
																						SCREAMING_SNAKE_CASE_   :	Any													=   measure_peak_memory_cpu(UpperCamelCase__       )
																						SCREAMING_SNAKE_CASE_   :	Any													=   Memory(UpperCamelCase__       ) if isinstance(UpperCamelCase__ ,    UpperCamelCase__       ) else memory_bytes
														if self.args.trace_memory_line_by_line:
																		SCREAMING_SNAKE_CASE_   :	Optional[int]													=   stop_memory_tracing(UpperCamelCase__       )
																		if memory is None:
																						SCREAMING_SNAKE_CASE_   :	str													=   summary.total
														else:
																		SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   None

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







import os

from bleurt import score  # From: git+https://github.com/google-research/bleurt.git

import datasets


A        =			datasets.logging.get_logger(__name__)


A        =			'\\n@inproceedings{bleurt,\n  title={BLEURT: Learning Robust Metrics for Text Generation},\n  author={Thibault Sellam and Dipanjan Das and Ankur P. Parikh},\n  booktitle={ACL},\n  year={2020},\n  url={https://arxiv.org/abs/2004.04696}\n}\n'

A        =			'\\nBLEURT a learnt evaluation metric for Natural Language Generation. It is built using multiple phases of transfer learning starting from a pretrained BERT model (Devlin et al. 2018)\nand then employing another pre-training phrase using synthetic data. Finally it is trained on WMT human annotations. You may run BLEURT out-of-the-box or fine-tune\nit for your specific application (the latter is expected to perform better).\n\nSee the project\'s README at https://github.com/google-research/bleurt#readme for more information.\n'

A        =			'\nBLEURT score.\n\nArgs:\n    `predictions` (list of str): prediction/candidate sentences\n    `references` (list of str): reference sentences\n    `checkpoint` BLEURT checkpoint. Will default to BLEURT-tiny if None.\n\nReturns:\n    \'scores\': List of scores.\nExamples:\n\n    >>> predictions = ["hello there", "general kenobi"]\n    >>> references = ["hello there", "general kenobi"]\n    >>> bleurt = datasets.load_metric("bleurt")\n    >>> results = bleurt.compute(predictions=predictions, references=references)\n    >>> print([round(v, 2) for v in results["scores"]])\n    [1.03, 1.04]\n'

A        =			{
    'bleurt-tiny-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-128.zip',
    'bleurt-tiny-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-tiny-512.zip',
    'bleurt-base-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-128.zip',
    'bleurt-base-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-base-512.zip',
    'bleurt-large-128': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-128.zip',
    'bleurt-large-512': 'https://storage.googleapis.com/bleurt-oss/bleurt-large-512.zip',
    'BLEURT-20-D3': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D3.zip',
    'BLEURT-20-D6': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D6.zip',
    'BLEURT-20-D12': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20-D12.zip',
    'BLEURT-20': 'https://storage.googleapis.com/bleurt-oss-21/BLEURT-20.zip',
}


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





		'''simple docstring'''



		def   __snake_case			(						self       ):
						return datasets.MetricInfo(
						    description=_DESCRIPTION ,    citation=_CITATION ,    homepage='https://github.com/google-research/bleurt' ,    inputs_description=_KWARGS_DESCRIPTION ,    features=datasets.Features(
						        {
						            'predictions': datasets.Value('string' ,    id='sequence'       ),
						            'references': datasets.Value('string' ,    id='sequence'       ),
						        }       ) ,    codebase_urls=['https://github.com/google-research/bleurt'] ,    reference_urls=['https://github.com/google-research/bleurt', 'https://arxiv.org/abs/2004.04696'] ,    )


		def   __snake_case			(						self ,    UpperCamelCase__       ):
						# check that config name specifies a valid BLEURT model
						if self.config_name == "default":
										logger.warning(
										    'Using default BLEURT-Base checkpoint for sequence maximum length 128. '
										    'You can use a bigger model for better results with e.g.: datasets.load_metric(\'bleurt\', \'bleurt-large-512\').'       )
										SCREAMING_SNAKE_CASE_   :	List[Any]													=   'bleurt-base-128'

						if self.config_name.lower() in CHECKPOINT_URLS:
										SCREAMING_SNAKE_CASE_   :	Tuple													=   self.config_name.lower()

						elif self.config_name.upper() in CHECKPOINT_URLS:
										SCREAMING_SNAKE_CASE_   :	Any													=   self.config_name.upper()

						else:
										raise KeyError(
										    F'''{self.config_name} model not found. You should supply the name of a model checkpoint for bleurt in {CHECKPOINT_URLS.keys()}'''       )

						# download the model checkpoint specified by self.config_name and set up the scorer
						SCREAMING_SNAKE_CASE_   :	List[Any]													=   dl_manager.download_and_extract(CHECKPOINT_URLS[checkpoint_name]       )
						SCREAMING_SNAKE_CASE_   :	int													=   score.BleurtScorer(os.path.join(UpperCamelCase__ ,    UpperCamelCase__       )       )






		def   __snake_case			(						self ,    UpperCamelCase__ ,    UpperCamelCase__       ):
						SCREAMING_SNAKE_CASE_   :	Tuple													=   self.scorer.score(references=UpperCamelCase__ ,    candidates=UpperCamelCase__       )
						return {"scores": scores}