


"""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

__magic_name__						: Tuple    =				logging.get_logger(__name__)





def       a_ (	lowercase__ :bool,			lowercase__ :bool     ):
			def run_func(lowercase__ :List[str]     ):
						@wraps(lowercase__     )
						def run_in_eager_mode(*lowercase__ :Any,			**lowercase__ :Tuple     ):
									return func(*lowercase__,			**lowercase__     )

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

						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       a_ (	lowercase__ :int,			lowercase__ :int,			lowercase__ :int     ):
			__lowerCamelCase          =      random.Random()
			__lowerCamelCase          =      [rng.randint(0,			vocab_size - 1     ) for i in range(batch_size * sequence_length     )]
			return tf.constant(lowercase__,			shape=(batch_size, sequence_length),			dtype=tf.intaa     )




class 			__snake_case   (lowerCamelCase  ):
							__a					 =       42
							__a					 =       42
							__a					 =       "TensorFlow"




							@property
							def   __a   (     self:    Optional[Any]   ):
										return tf.__version__




							def   __a   (     self:    Union[str, Any]     ,							A_:    str     ,							A_:    int     ,							A_:    int   ):
										# initialize GPU on separate process
										__lowerCamelCase          =      self.args.strategy
										if strategy is None:
													raise ValueError("""A device strategy has to be initialized before using TensorFlow."""   )
										__lowerCamelCase          =      self._prepare_inference_func(A_     ,							A_     ,							A_   )
										return self._measure_speed(_inference   )




							def   __a   (     self:    Union[str, Any]     ,							A_:    str     ,							A_:    int     ,							A_:    int   ):
										__lowerCamelCase          =      self.args.strategy
										if strategy is None:
													raise ValueError("""A device strategy has to be initialized before using TensorFlow."""   )
										__lowerCamelCase          =      self._prepare_train_func(A_     ,							A_     ,							A_   )
										return self._measure_speed(_train   )




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




							def   __a   (     self:    Optional[Any]     ,							A_:    str     ,							A_:    int     ,							A_:    int   ):
										if self.args.is_gpu:
													tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]     ,							A_   )
										__lowerCamelCase          =      self.args.strategy
										if strategy is None:
													raise ValueError("""A device strategy has to be initialized before using TensorFlow."""   )

										__lowerCamelCase          =      self._prepare_train_func(A_     ,							A_     ,							A_   )
										return self._measure_memory(_train   )




							def   __a   (     self:    Optional[Any]     ,							A_:    str     ,							A_:    int     ,							A_:    int   ):
										__lowerCamelCase          =      self.config_dict[model_name]

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

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

										# encoder-decoder has vocab size saved differently
										__lowerCamelCase          =      config.vocab_size if hasattr(A_     ,							"""vocab_size"""   ) else config.encoder.vocab_size
										__lowerCamelCase          =      random_input_ids(A_     ,							A_     ,							A_   )

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

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

										__lowerCamelCase          =      encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

										return _inference




							def   __a   (     self:    Union[str, Any]     ,							A_:    str     ,							A_:    int     ,							A_:    int   ):
										__lowerCamelCase          =      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."""   )

										__lowerCamelCase          =      (
										    hasattr(A_     ,							"""architectures"""   )
										    and isinstance(config.architectures     ,							A_   )
										    and len(config.architectures   ) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
													try:
																__lowerCamelCase          =      """TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
																__lowerCamelCase          =      __import__("""transformers"""     ,							fromlist=[model_class]   )
																__lowerCamelCase          =      getattr(A_     ,							A_   )
																__lowerCamelCase          =      model_cls(A_   )
													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:
													__lowerCamelCase          =      TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](A_   )

										# encoder-decoder has vocab size saved differently
										__lowerCamelCase          =      config.vocab_size if hasattr(A_     ,							"""vocab_size"""   ) else config.encoder.vocab_size
										__lowerCamelCase          =      random_input_ids(A_     ,							A_     ,							A_   )

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

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

										__lowerCamelCase          =      encoder_decoder_train if config.is_encoder_decoder else encoder_train

										return _train




							def   __a   (     self:    Optional[int]     ,							A_:    Any   ):
										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(A_     ,							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
																__lowerCamelCase          =      timeit.repeat(
																    A_     ,							repeat=self.args.repeat     ,							number=10     ,							)

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




							def   __a   (     self:    int     ,							A_:    Callable[[], None]   ):
										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."""   )
																			__lowerCamelCase          =      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."""   )
																						__lowerCamelCase          =      """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()
																						__lowerCamelCase          =      nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx   )
																						__lowerCamelCase          =      nvml.nvmlDeviceGetMemoryInfo(A_   )
																						__lowerCamelCase          =      meminfo.used
																						__lowerCamelCase          =      Memory(A_   )
																						# 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."""   )
																						__lowerCamelCase          =      None
																			else:
																						__lowerCamelCase          =      measure_peak_memory_cpu(A_   )
																						__lowerCamelCase          =      Memory(A_   ) if isinstance(A_     ,							A_   ) else memory_bytes
																if self.args.trace_memory_line_by_line:
																			__lowerCamelCase          =      stop_memory_tracing(A_   )
																			if memory is None:
																						__lowerCamelCase          =      summary.total
																else:
																			__lowerCamelCase          =      None

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








"""simple docstring"""




__magic_name__						: Optional[Any]    =				'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'





def       a_ (	lowercase__ :bytes     ):
			# Make sure the supplied data is a bytes-like object
			if not isinstance(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      f'a bytes-like object is required, not \'{data.__class__.__name__}\''
						raise TypeError(lowercase__     )

			__lowerCamelCase          =      """""".join(bin(lowercase__     )[2:].zfill(8     ) for byte in data     )

			__lowerCamelCase          =      len(lowercase__     ) % 6 != 0

			if padding_needed:
						# The padding that will be added later
						__lowerCamelCase          =      B"""=""" * ((6 - len(lowercase__     ) % 6) // 2)

						# Append binary_stream with arbitrary binary digits (0's by default) to make its
						# length a multiple of 6.
						binary_stream += "0" * (6 - len(lowercase__     ) % 6)
			else:
						__lowerCamelCase          =      B""""""

			# Encode every 6 binary digits to their corresponding Base64 character
			return (
			    "".join(
			        B64_CHARSET[int(binary_stream[index : index + 6],			2     )]
			        for index in range(0,			len(lowercase__     ),			6     )     ).encode()
			    + padding
			)





def       a_ (	lowercase__ :str     ):
			# Make sure encoded_data is either a string or a bytes-like object
			if not isinstance(lowercase__,			lowercase__     ) and not isinstance(lowercase__,			lowercase__     ):
						__lowerCamelCase          =      (
						    """argument should be a bytes-like object or ASCII string, """
						    f'not \'{encoded_data.__class__.__name__}\''
						)
						raise TypeError(lowercase__     )

			# In case encoded_data is a bytes-like object, make sure it contains only
			# ASCII characters so we convert it to a string object
			if isinstance(lowercase__,			lowercase__     ):
						try:
									__lowerCamelCase          =      encoded_data.decode("""utf-8"""     )
						except UnicodeDecodeError:
									raise ValueError("""base64 encoded data should only contain ASCII characters"""     )

			__lowerCamelCase          =      encoded_data.count("""="""     )

			# Check if the encoded string contains non base64 characters
			if padding:
						assert all(
						    char in B64_CHARSET for char in encoded_data[:-padding]     ), "Invalid base64 character(s) found."
			else:
						assert all(
						    char in B64_CHARSET for char in encoded_data     ), "Invalid base64 character(s) found."

			# Check the padding
			assert len(lowercase__     ) % 4 == 0 and padding < 3, "Incorrect padding"

			if padding:
						# Remove padding if there is one
						__lowerCamelCase          =      encoded_data[:-padding]

						__lowerCamelCase          =      """""".join(
						    bin(B64_CHARSET.index(lowercase__     )     )[2:].zfill(6     ) for char in encoded_data     )[: -padding * 2]
			else:
						__lowerCamelCase          =      """""".join(
						    bin(B64_CHARSET.index(lowercase__     )     )[2:].zfill(6     ) for char in encoded_data     )

			__lowerCamelCase          =      [
			    int(binary_stream[index : index + 8],			2     )
			    for index in range(0,			len(lowercase__     ),			8     )
			]

			return bytes(lowercase__     )


if __name__ == "__main__":
					import doctest

					doctest.testmod()




