



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_							:      Any      =  logging.get_logger(__name__)
def  UpperCamelCase						(       _A			:				bool			,							_A			:				bool    )->       int:





						"""simple docstring"""





						def run_func(_A			:				Tuple    ):
												@wraps(_A    )
												def run_in_eager_mode(*_A			:				List[str]			,							**_A			:				List[str]    ):
																		return func(*_A			,							**_A    )

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

												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						(       _A			:				int			,							_A			:				int			,							_A			:				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(_A			,							shape=(batch_size, sequence_length)			,							dtype=tf.intaa    )






class      UpperCamelCase    (						_UpperCAmelCase     ):
	lowerCAmelCase	:   TensorFlowBenchmarkArguments
	lowerCAmelCase	:   PretrainedConfig
	lowerCAmelCase	:   str				=	"TensorFlow"



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



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							# 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(UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		)
							return self._measure_speed(_inference		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							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(UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		)
							return self._measure_speed(_train		)



	def 			__A							(       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__		)
							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(UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		)
							return self._measure_memory(_inference		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							if self.args.is_gpu:
													tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]			,	UpperCAmelCase__		)
							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(UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		)
							return self._measure_memory(_train		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       self.config_dict[model_name]

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

							A__							=       (
							    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:
																			A__							=       "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
																			A__							=       __import__("transformers"			,	fromlist=[model_class]		)
																			A__							=       getattr(UpperCAmelCase__			,	UpperCAmelCase__		)
																			A__							=       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:
													A__							=       TF_MODEL_MAPPING[config.__class__](UpperCAmelCase__		)

							# encoder-decoder has vocab size saved differently
							A__							=       config.vocab_size if hasattr(UpperCAmelCase__			,	"vocab_size"		) else config.encoder.vocab_size
							A__							=       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__		)

							A__							=       encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

							return _inference



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							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(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:
																			A__							=       "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
																			A__							=       __import__("transformers"			,	fromlist=[model_class]		)
																			A__							=       getattr(UpperCAmelCase__			,	UpperCAmelCase__		)
																			A__							=       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:
													A__							=       TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](UpperCAmelCase__		)

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

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

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

							A__							=       encoder_decoder_train if config.is_encoder_decoder else encoder_train

							return _train



	def 			__A							(       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
																			A__							=       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 			__A							(       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."		)
																									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(UpperCAmelCase__		)
																															A__							=       meminfo.used
																															A__							=       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."		)
																															A__							=       None
																									else:
																															A__							=       measure_peak_memory_cpu(UpperCAmelCase__		)
																															A__							=       Memory(UpperCAmelCase__		) if isinstance(UpperCAmelCase__			,	UpperCAmelCase__		) else memory_bytes
																			if self.args.trace_memory_line_by_line:
																									A__							=       stop_memory_tracing(UpperCAmelCase__		)
																									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










from manim import *






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def 			__A							(       self		):
							A__							=       Rectangle(height=0.5			,	width=0.5		)
							A__							=       Rectangle(height=0.46			,	width=0.46		).set_stroke(width=0		)

							A__							=       [mem.copy() for i in range(6		)]
							A__							=       [mem.copy() for i in range(6		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       VGroup(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       Text("CPU"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0.5			,	aligned_edge=UpperCAmelCase__		)
							cpu.move_to([-2.5, -0.5, 0]		)
							self.add(UpperCAmelCase__		)

							A__							=       [mem.copy() for i in range(4		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       Text("GPU"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0.5			,	aligned_edge=UpperCAmelCase__		)
							gpu.move_to([-1, -1, 0]		)
							self.add(UpperCAmelCase__		)

							A__							=       [mem.copy() for i in range(6		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)

							A__							=       Text("Model"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0.5			,	aligned_edge=UpperCAmelCase__		)
							model.move_to([3, -1.0, 0]		)
							self.add(UpperCAmelCase__		)

							A__							=       []
							for i, rect in enumerate(UpperCAmelCase__		):
													rect.set_stroke(UpperCAmelCase__		)
													# target = fill.copy().set_fill(YELLOW, opacity=0.7)
													# target.move_to(rect)
													# self.add(target)

													A__							=       Rectangle(height=0.46 / 4			,	width=0.46 / 3		).set_stroke(width=0.0		).set_fill(UpperCAmelCase__			,	opacity=0.7		)

													if i == 0:
																			cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT		)			,	buff=0.02			,	direction=UpperCAmelCase__		)
																			cpu_target.set_x(cpu_target.get_x() + 0.1		)
													elif i == 3:
																			cpu_target.next_to(cpu_targs[0]			,	direction=UpperCAmelCase__			,	buff=0.0		)
													else:
																			cpu_target.next_to(cpu_targs[i - 1]			,	direction=UpperCAmelCase__			,	buff=0.0		)
													self.add(UpperCAmelCase__		)
													cpu_targs.append(UpperCAmelCase__		)

							A__							=       [mem.copy() for i in range(6		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)

							A__							=       Text("Loaded Checkpoint"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	aligned_edge=UpperCAmelCase__			,	buff=0.4		)
							checkpoint.move_to([3, 0.5, 0]		)

							A__							=       Square(side_length=2.2		)
							key.move_to([-5, 2, 0]		)

							A__							=       MarkupText(
							    F"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""			,	font_size=18			,	)

							key_text.move_to([-5, 2.4, 0]		)

							self.add(UpperCAmelCase__			,	UpperCAmelCase__		)

							A__							=       MarkupText(
							    F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""			,	font_size=18			,	)

							blue_text.next_to(UpperCAmelCase__			,	DOWN * 2.4			,	aligned_edge=key_text.get_left()		)

							A__							=       MarkupText(
							    F"""Next, a <i><span fgcolor=\"{BLUE}\">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor=\"{BLUE}\">single shard</span>."""			,	font_size=24			,	)
							step_a.move_to([2, 2, 0]		)
							self.play(Write(UpperCAmelCase__		)			,	Write(UpperCAmelCase__		)		)

							self.play(Write(UpperCAmelCase__			,	run_time=1		)			,	Create(UpperCAmelCase__			,	run_time=1		)		)

							A__							=       []
							A__							=       []
							for i, rect in enumerate(UpperCAmelCase__		):
													A__							=       fill.copy().set_fill(UpperCAmelCase__			,	opacity=0.7		)
													target.move_to(UpperCAmelCase__		)
													first_animations.append(GrowFromCenter(UpperCAmelCase__			,	run_time=1		)		)

													A__							=       target.copy()
													cpu_target.generate_target()
													if i < 5:
																			cpu_target.target.move_to(cpu_left_col_base[i + 1]		)
													else:
																			cpu_target.target.move_to(cpu_right_col_base[i - 5]		)
													second_animations.append(MoveToTarget(UpperCAmelCase__			,	run_time=1.5		)		)

							self.play(*UpperCAmelCase__		)
							self.play(*UpperCAmelCase__		)
							self.wait()





