




'''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_    	=				logging.get_logger(__name__)




def   lowercase__(			__UpperCamelCase:					int      ,__UpperCamelCase:					List[str]    ):

      """simple docstring"""






      def run_func(__UpperCamelCase:					Tuple    ):
            @wraps(__UpperCamelCase    )
            def run_in_eager_mode(*__UpperCamelCase:					Dict      ,**__UpperCamelCase:					Any    ):
                  return func(*__UpperCamelCase      ,**__UpperCamelCase    )

            @wraps(__UpperCamelCase    )
            @tf.function(experimental_compile=__UpperCamelCase    )
            def run_in_graph_mode(*__UpperCamelCase:					Dict      ,**__UpperCamelCase:					Optional[Any]    ):
                  return func(*__UpperCamelCase      ,**__UpperCamelCase    )

            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   lowercase__(			__UpperCamelCase:					str      ,__UpperCamelCase:					Dict      ,__UpperCamelCase:					int    ):

      """simple docstring"""






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






class 							_a				(			lowerCamelCase__ ):





       '''simple docstring'''






       A						:				TensorFlowBenchmarkArguments
       A						:				PretrainedConfig
       A						:				str        =						"TensorFlow"







       @property
       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             return tf.__version__







       def        UpperCamelCase_							(      self,							A,							A,							A  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Dict    					=						self.args.strategy
             if strategy is None:
                   raise ValueError('A device strategy has to be initialized before using TensorFlow.'  )
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						self._prepare_inference_func(snake_case__,							snake_case__,							snake_case__  )
             return self._measure_speed(_inference  )







       def        UpperCamelCase_							(      self,							A,							A,							A  ):







             '''simple docstring'''







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







       def        UpperCamelCase_							(      self,							A,							A,							A  ):







             '''simple docstring'''







             if self.args.is_gpu:
                   tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx],							snake_case__  )
             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   :      List[Any]    					=						self._prepare_inference_func(snake_case__,							snake_case__,							snake_case__  )
             return self._measure_memory(_inference  )







       def        UpperCamelCase_							(      self,							A,							A,							A  ):







             '''simple docstring'''







             if self.args.is_gpu:
                   tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx],							snake_case__  )
             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   :      Optional[int]    					=						self._prepare_train_func(snake_case__,							snake_case__,							snake_case__  )
             return self._measure_memory(_train  )







       def        UpperCamelCase_							(      self,							A,							A,							A  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Dict    					=						self.config_dict[model_name]

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

             SCREAMING_SNAKE_CASE   :      str    					=						(
                 hasattr(snake_case__,							'architectures'  )
                 and isinstance(config.architectures,							snake_case__  )
                 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   :      Any    					=						__import__('transformers',							fromlist=[model_class]  )
                         SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						getattr(snake_case__,							snake_case__  )
                         SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						model_cls(snake_case__  )
                   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   :      Union[str, Any]    					=						TF_MODEL_MAPPING[config.__class__](snake_case__  )

             # encoder-decoder has vocab size saved differently
             SCREAMING_SNAKE_CASE   :      Tuple    					=						config.vocab_size if hasattr(snake_case__,							'vocab_size'  ) else config.encoder.vocab_size
             SCREAMING_SNAKE_CASE   :      List[str]    					=						random_input_ids(snake_case__,							snake_case__,							snake_case__  )

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

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

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

             return _inference







       def        UpperCamelCase_							(      self,							A,							A,							A  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						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   :      Optional[Any]    					=						(
                 hasattr(snake_case__,							'architectures'  )
                 and isinstance(config.architectures,							snake_case__  )
                 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   :      Optional[Any]    					=						__import__('transformers',							fromlist=[model_class]  )
                         SCREAMING_SNAKE_CASE   :      List[Any]    					=						getattr(snake_case__,							snake_case__  )
                         SCREAMING_SNAKE_CASE   :      Dict    					=						model_cls(snake_case__  )
                   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   :      Tuple    					=						TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case__  )

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

             @run_with_tf_optimizations(self.args.eager_mode,							self.args.use_xla  )
             def encoder_decoder_train():
                   SCREAMING_SNAKE_CASE   :      Any    					=						model(snake_case__,							decoder_input_ids=snake_case__,							labels=snake_case__,							training=snake_case__  )[0]
                   SCREAMING_SNAKE_CASE   :      List[Any]    					=						tf.gradients(snake_case__,							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(snake_case__,							labels=snake_case__,							training=snake_case__  )[0]
                   SCREAMING_SNAKE_CASE   :      Any    					=						tf.gradients(snake_case__,							model.trainable_variables  )
                   return gradients

             SCREAMING_SNAKE_CASE   :      str    					=						encoder_decoder_train if config.is_encoder_decoder else encoder_train

             return _train







       def        UpperCamelCase_							(      self,							A  ):







             '''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(snake_case__,							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(
                             snake_case__,							repeat=self.args.repeat,							number=10,							)

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







       def        UpperCamelCase_							(      self,							A  ):







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

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









import tensorflow as tf

from ...tf_utils import shape_list
class     UpperCamelCase__					(     tf.keras.layers.Layer				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__=1,   snake_case__=False,   **snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									super().__init__(**snake_case__  )

									lowercase_	:   Tuple  				=     vocab_size
									lowercase_	:   Union[str, Any]  				=     d_embed
									lowercase_	:   Optional[Any]  				=     d_proj

									lowercase_	:   Optional[Any]  				=     cutoffs + [vocab_size]
									lowercase_	:   Optional[int]  				=     [0] + self.cutoffs
									lowercase_	:   Union[str, Any]  				=     div_val

									lowercase_	:   Union[str, Any]  				=     self.cutoffs[0]
									lowercase_	:   List[str]  				=     len(self.cutoffs  ) - 1
									lowercase_	:   List[str]  				=     self.shortlist_size + self.n_clusters
									lowercase_	:   Tuple  				=     keep_order

									lowercase_	:   Dict  				=     []
									lowercase_	:   int  				=     []






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									if self.n_clusters > 0:
																lowercase_	:   List[str]  				=     self.add_weight(
																    shape=(self.n_clusters, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_weight"""  )
																lowercase_	:   List[Any]  				=     self.add_weight(
																    shape=(self.n_clusters,),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_bias"""  )

									if self.div_val == 1:
																for i in range(len(self.cutoffs  )  ):
																							if self.d_proj != self.d_embed:
																														lowercase_	:   Any  				=     self.add_weight(
																														    shape=(self.d_embed, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}""",   )
																														self.out_projs.append(snake_case__  )
																							else:
																														self.out_projs.append(snake_case__  )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									else:
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[Any]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							lowercase_	:   Union[str, Any]  				=     self.d_embed // (self.div_val**i)

																							lowercase_	:   int  				=     self.add_weight(
																							    shape=(d_emb_i, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}"""  )
																							self.out_projs.append(snake_case__  )
																							lowercase_	:   Any  				=     self.add_weight(
																							    shape=(r_idx - l_idx, d_emb_i),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   Tuple  				=     self.add_weight(
																							    shape=(r_idx - l_idx,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									super().build(snake_case__  )






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__,   snake_case__,   snake_case__=None  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Dict  				=     x
									if proj is not None:
																lowercase_	:   List[Any]  				=     tf.einsum("""ibd,ed->ibe""",   snake_case__,   snake_case__  )
									return tf.einsum("""ibd,nd->ibn""",   snake_case__,   snake_case__  ) + b






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     shape_list(snake_case__  )
									lowercase_	:   Optional[Any]  				=     tf.range(lp_size[0],   dtype=target.dtype  )
									lowercase_	:   str  				=     tf.stack([r, target],   1  )
									return tf.gather_nd(snake_case__,   snake_case__  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__,   snake_case__=True,   snake_case__=False  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   Any  				=     0
									if self.n_clusters == 0:
																lowercase_	:   Union[str, Any]  				=     self._logit(snake_case__,   self.out_layers[0][0],   self.out_layers[0][1],   self.out_projs[0]  )
																if target is not None:
																							lowercase_	:   str  				=     tf.nn.sparse_softmax_cross_entropy_with_logits(labels=snake_case__,   logits=snake_case__  )
																lowercase_	:   Optional[int]  				=     tf.nn.log_softmax(snake_case__,   axis=-1  )
									else:
																lowercase_	:   int  				=     shape_list(snake_case__  )
																lowercase_	:   List[Any]  				=     []
																lowercase_	:   Tuple  				=     tf.zeros(hidden_sizes[:2]  )
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[int]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							if target is not None:
																														lowercase_	:   Tuple  				=     (target >= l_idx) & (target < r_idx)
																														lowercase_	:   Union[str, Any]  				=     tf.where(snake_case__  )
																														lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  ) - l_idx

																							if self.div_val == 1:
																														lowercase_	:   Optional[Any]  				=     self.out_layers[0][0][l_idx:r_idx]
																														lowercase_	:   Union[str, Any]  				=     self.out_layers[0][1][l_idx:r_idx]
																							else:
																														lowercase_	:   Dict  				=     self.out_layers[i][0]
																														lowercase_	:   int  				=     self.out_layers[i][1]

																							if i == 0:
																														lowercase_	:   Optional[int]  				=     tf.concat([cur_W, self.cluster_weight],   0  )
																														lowercase_	:   Optional[Any]  				=     tf.concat([cur_b, self.cluster_bias],   0  )

																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[0]  )
																														lowercase_	:   List[str]  				=     tf.nn.log_softmax(snake_case__  )
																														out.append(head_logprob[..., : self.cutoffs[0]]  )
																														if target is not None:
																																					lowercase_	:   int  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																							else:
																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[i]  )
																														lowercase_	:   Dict  				=     tf.nn.log_softmax(snake_case__  )
																														lowercase_	:   Optional[int]  				=     self.cutoffs[0] + i - 1  # No probability for the head cluster
																														lowercase_	:   int  				=     head_logprob[..., cluster_prob_idx, None] + tail_logprob
																														out.append(snake_case__  )
																														if target is not None:
																																					lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Union[str, Any]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																																					cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
																							if target is not None:
																														loss += tf.scatter_nd(snake_case__,   -cur_logprob,   shape_list(snake_case__  )  )
																lowercase_	:   List[str]  				=     tf.concat(snake_case__,   axis=-1  )

									if target is not None:
																if return_mean:
																							lowercase_	:   Tuple  				=     tf.reduce_mean(snake_case__  )
																# Add the training-time loss value to the layer using `self.add_loss()`.
																self.add_loss(snake_case__  )

																# Log the loss as a metric (we could log arbitrary metrics,
																# including different metrics for training and inference.
																self.add_metric(snake_case__,   name=self.name,   aggregation="""mean""" if return_mean else """"""  )

									return out