




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

a_	   =    logging.get_logger(__name__)







def 			__UpperCAmelCase			(      __UpperCamelCase     ,							__UpperCamelCase     ):





			def run_func(__UpperCamelCase     ):
						@wraps(_UpperCAmelCase     )
						def run_in_eager_mode(*__UpperCamelCase     ,							**__UpperCamelCase     ):
									return func(*_UpperCAmelCase     ,							**_UpperCAmelCase     )

						@wraps(_UpperCAmelCase     )
						@tf.function(experimental_compile=_UpperCAmelCase     )
						def run_in_graph_mode(*__UpperCamelCase     ,							**__UpperCamelCase     ):
									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 			__UpperCAmelCase			(      __UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ):
			__lowercase       :   Union[str, Any]        =    random.Random()
			__lowercase       :   List[str]        =    [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    UpperCAmelCase_  (       UpperCAmelCase__						):
		UpperCamelCase						 =42
		UpperCamelCase						 =42
		UpperCamelCase						 ="TensorFlow"
		@property
		def      _lowerCamelCase				(		self    )		->    int:
					return tf.__version__
		def      _lowerCamelCase				(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_    )		->    Dict:
					__lowercase       :   str        =    self.args.strategy
					if strategy is None:
								raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''    )
					__lowercase       :   int        =    self._prepare_inference_func(snake_case_							,		snake_case_							,		snake_case_    )
					return self._measure_speed(_inference    )
		def      _lowerCamelCase				(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_    )		->    Union[str, Any]:
					__lowercase       :   Any        =    self.args.strategy
					if strategy is None:
								raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''    )
					__lowercase       :   str        =    self._prepare_train_func(snake_case_							,		snake_case_							,		snake_case_    )
					return self._measure_speed(_train    )
		def      _lowerCamelCase				(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_    )		->    Any:

					if self.args.is_gpu:
								tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]							,		snake_case_    )
					__lowercase       :   Any        =    self.args.strategy
					if strategy is None:
								raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''    )
					__lowercase       :   Any        =    self._prepare_inference_func(snake_case_							,		snake_case_							,		snake_case_    )
					return self._measure_memory(_inference    )
		def      _lowerCamelCase				(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_    )		->    Optional[Any]:

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

					__lowercase       :   str        =    self._prepare_train_func(snake_case_							,		snake_case_							,		snake_case_    )
					return self._measure_memory(_train    )
		def      _lowerCamelCase				(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_    )		->    Union[str, Any]:
					__lowercase       :   str        =    self.config_dict[model_name]

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

					__lowercase       :   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:
											__lowercase       :   str        =    'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
											__lowercase       :   Optional[Any]        =    __import__('''transformers'''							,		fromlist=[model_class]    )
											__lowercase       :   int        =    getattr(snake_case_							,		snake_case_    )
											__lowercase       :   List[str]        =    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:
								__lowercase       :   List[str]        =    TF_MODEL_MAPPING[config.__class__](snake_case_    )

					# encoder-decoder has vocab size saved differently
					__lowercase       :   Optional[int]        =    config.vocab_size if hasattr(snake_case_							,		'''vocab_size'''    ) else config.encoder.vocab_size
					__lowercase       :   Optional[int]        =    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_    )

					__lowercase       :   str        =    encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

					return _inference
		def      _lowerCamelCase				(		self							,		UpperCamelCase_							,		UpperCamelCase_							,		UpperCamelCase_    )		->    Union[str, Any]:
					__lowercase       :   List[Any]        =    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.'''    )

					__lowercase       :   int        =    (
					    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:
											__lowercase       :   Any        =    'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
											__lowercase       :   Optional[Any]        =    __import__('''transformers'''							,		fromlist=[model_class]    )
											__lowercase       :   Tuple        =    getattr(snake_case_							,		snake_case_    )
											__lowercase       :   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:
								__lowercase       :   Any        =    TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case_    )

					# encoder-decoder has vocab size saved differently
					__lowercase       :   Optional[int]        =    config.vocab_size if hasattr(snake_case_							,		'''vocab_size'''    ) else config.encoder.vocab_size
					__lowercase       :   Dict        =    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():
								__lowercase       :   int        =    model(snake_case_							,		decoder_input_ids=snake_case_							,		labels=snake_case_							,		training=snake_case_    )[0]
								__lowercase       :   Union[str, 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():
								__lowercase       :   Any        =    model(snake_case_							,		labels=snake_case_							,		training=snake_case_    )[0]
								__lowercase       :   List[Any]        =    tf.gradients(snake_case_							,		model.trainable_variables    )
								return gradients

					__lowercase       :   Tuple        =    encoder_decoder_train if config.is_encoder_decoder else encoder_train

					return _train
		def      _lowerCamelCase				(		self							,		UpperCamelCase_    )		->    int:

					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
											__lowercase       :   List[str]        =    timeit.repeat(
											    snake_case_							,		repeat=self.args.repeat							,		number=10							,		)

											return min(snake_case_    ) / 1_0.0
								except ResourceExhaustedError as e:
											self.print_fn(F"""Doesn't fit on GPU. {e}"""    )
		def      _lowerCamelCase				(		self							,		UpperCamelCase_    )		->    List[Any]:
					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.'''    )
														__lowercase       :   Optional[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.'''    )
																	__lowercase       :   Tuple        =    '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()
																	__lowercase       :   List[str]        =    nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx    )
																	__lowercase       :   int        =    nvml.nvmlDeviceGetMemoryInfo(snake_case_    )
																	__lowercase       :   int        =    meminfo.used
																	__lowercase       :   int        =    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.'''    )
																	__lowercase       :   Any        =    None
														else:
																	__lowercase       :   List[str]        =    measure_peak_memory_cpu(snake_case_    )
																	__lowercase       :   Any        =    Memory(snake_case_    ) if isinstance(snake_case_							,		snake_case_    ) else memory_bytes
											if self.args.trace_memory_line_by_line:
														__lowercase       :   List[str]        =    stop_memory_tracing(snake_case_    )
														if memory is None:
																	__lowercase       :   Union[str, Any]        =    summary.total
											else:
														__lowercase       :   Any        =    None

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










"""simple docstring"""




import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


lowerCamelCase_     :						Tuple						=					logging.get_logger(__name__)





def 						UpperCAmelCase__		(     _UpperCAmelCase					):






  """simple docstring"""





  A_  :       Optional[Any]										=  R'\w+[.]\d+'
  A_  :       int										=  re.findall(_UpperCAmelCase ,					_UpperCAmelCase					)
  for pat in pats:
    A_  :       Optional[int]										=  key.replace(_UpperCAmelCase ,					'_'.join(pat.split('.'					)					)					)
  return key





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					):






  """simple docstring"""





  A_  :       List[Any]										=  pt_tuple_key[:-1] + ('scale',)
  if (
      any('norm' in str_ for str_ in pt_tuple_key					)
      and (pt_tuple_key[-1] == "bias")
      and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
      and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
  ):
    A_  :       Union[str, Any]										=  pt_tuple_key[:-1] + ('scale',)
    return renamed_pt_tuple_key, pt_tensor
  elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
    A_  :       List[str]										=  pt_tuple_key[:-1] + ('scale',)
    return renamed_pt_tuple_key, pt_tensor

  # embedding
  if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
    A_  :       Optional[Any]										=  pt_tuple_key[:-1] + ('embedding',)
    return renamed_pt_tuple_key, pt_tensor

  # conv layer
  A_  :       int										=  pt_tuple_key[:-1] + ('kernel',)
  if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
    A_  :       str										=  pt_tensor.transpose(2 ,					3 ,					1 ,					0					)
    return renamed_pt_tuple_key, pt_tensor

  # linear layer
  A_  :       Optional[Any]										=  pt_tuple_key[:-1] + ('kernel',)
  if pt_tuple_key[-1] == "weight":
    A_  :       Optional[Any]										=  pt_tensor.T
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm weight
  A_  :       Tuple										=  pt_tuple_key[:-1] + ('weight',)
  if pt_tuple_key[-1] == "gamma":
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm bias
  A_  :       Optional[int]										=  pt_tuple_key[:-1] + ('bias',)
  if pt_tuple_key[-1] == "beta":
    return renamed_pt_tuple_key, pt_tensor

  return pt_tuple_key, pt_tensor





def 						UpperCAmelCase__		(     _UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase=42					):






  """simple docstring"""





  A_  :       int										=  {k: v.numpy() for k, v in pt_state_dict.items()}

  # Step 2: Since the model is stateless, get random Flax params
  A_  :       Union[str, Any]										=  flax_model.init_weights(PRNGKey(_UpperCAmelCase					)					)

  A_  :       Optional[Any]										=  flatten_dict(_UpperCAmelCase					)
  A_  :       Tuple										=  {}

  # Need to change some parameters name to match Flax names
  for pt_key, pt_tensor in pt_state_dict.items():
    A_  :       Any										=  rename_key(_UpperCAmelCase					)
    A_  :       List[str]										=  tuple(renamed_pt_key.split('.'					)					)

    # Correctly rename weight parameters
    A_						,       A_  :       Union[str, Any]										=  rename_key_and_reshape_tensor(_UpperCAmelCase ,					_UpperCAmelCase ,					_UpperCAmelCase					)

    if flax_key in random_flax_state_dict:
      if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
        raise ValueError(
            f"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
            f"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""					)

        # also add unexpected weight so that warning is thrown
    A_  :       str										=  jnp.asarray(_UpperCAmelCase					)

  return unflatten_dict(_UpperCAmelCase					)