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

_SCREAMING_SNAKE_CASE	       =      logging.get_logger(__name__)
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Union[str, Any]:







							"""simple docstring"""


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

														@wraps(SCREAMING_SNAKE_CASE )
														@tf.function(experimental_compile=SCREAMING_SNAKE_CASE )
														def run_in_graph_mode(*SCREAMING_SNAKE_CASE							,    **SCREAMING_SNAKE_CASE ):
																					return func(*SCREAMING_SNAKE_CASE							,    **SCREAMING_SNAKE_CASE )

														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   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     ["tf.Tensor"]:







							"""simple docstring"""


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







class 							__magic_name__     (		lowercase__   ):
					_SCREAMING_SNAKE_CASE       :					TensorFlowBenchmarkArguments
					_SCREAMING_SNAKE_CASE       :					PretrainedConfig
					_SCREAMING_SNAKE_CASE       :					str     							=   "TensorFlow"





					@property
					def 							lowerCAmelCase      (				self   :							int    ):
												return tf.__version__





					def 							lowerCAmelCase      (				self   :							str							,						snake_case_   :							str							,						snake_case_   :							int							,						snake_case_   :							int    ):
												# initialize GPU on separate process
												__snake_case     			=       self.args.strategy
												if strategy is None:
																			raise ValueError("A device strategy has to be initialized before using TensorFlow."    )
												__snake_case     			=       self._prepare_inference_func(snake_case_							,						snake_case_							,						snake_case_    )
												return self._measure_speed(_inference    )





					def 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							str							,						snake_case_   :							int							,						snake_case_   :							int    ):
												__snake_case     			=       self.args.strategy
												if strategy is None:
																			raise ValueError("A device strategy has to be initialized before using TensorFlow."    )
												__snake_case     			=       self._prepare_train_func(snake_case_							,						snake_case_							,						snake_case_    )
												return self._measure_speed(_train    )





					def 							lowerCAmelCase      (				self   :							str							,						snake_case_   :							str							,						snake_case_   :							int							,						snake_case_   :							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]							,						snake_case_    )
												__snake_case     			=       self.args.strategy
												if strategy is None:
																			raise ValueError("A device strategy has to be initialized before using TensorFlow."    )
												__snake_case     			=       self._prepare_inference_func(snake_case_							,						snake_case_							,						snake_case_    )
												return self._measure_memory(_inference    )





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

												__snake_case     			=       self._prepare_train_func(snake_case_							,						snake_case_							,						snake_case_    )
												return self._measure_memory(_train    )





					def 							lowerCAmelCase      (				self   :							Dict							,						snake_case_   :							str							,						snake_case_   :							int							,						snake_case_   :							int    ):
												__snake_case     			=       self.config_dict[model_name]

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

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

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

												__snake_case     			=       encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

												return _inference





					def 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							str							,						snake_case_   :							int							,						snake_case_   :							int    ):
												__snake_case     			=       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."    )

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

												# encoder-decoder has vocab size saved differently
												__snake_case     			=       config.vocab_size if hasattr(snake_case_							,						"vocab_size"    ) else config.encoder.vocab_size
												__snake_case     			=       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():
																			__snake_case     			=       model(snake_case_							,						decoder_input_ids=snake_case_							,						labels=snake_case_							,						training=snake_case_    )[0]
																			__snake_case     			=       tf.gradients(snake_case_							,						model.trainable_variables    )
																			return gradients

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

												__snake_case     			=       encoder_decoder_train if config.is_encoder_decoder else encoder_train

												return _train





					def 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							Optional[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(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
																										__snake_case     			=       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 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							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."    )
																																	__snake_case     			=       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."    )
																																								__snake_case     			=       "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()
																																								__snake_case     			=       nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx    )
																																								__snake_case     			=       nvml.nvmlDeviceGetMemoryInfo(snake_case_    )
																																								__snake_case     			=       meminfo.used
																																								__snake_case     			=       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."    )
																																								__snake_case     			=       None
																																	else:
																																								__snake_case     			=       measure_peak_memory_cpu(snake_case_    )
																																								__snake_case     			=       Memory(snake_case_    ) if isinstance(snake_case_							,						snake_case_    ) else memory_bytes
																										if self.args.trace_memory_line_by_line:
																																	__snake_case     			=       stop_memory_tracing(snake_case_    )
																																	if memory is None:
																																								__snake_case     			=       summary.total
																										else:
																																	__snake_case     			=       None

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


"""simple docstring"""







from pathlib import Path
from typing import List

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import get_tests_dir, is_tool_test
from transformers.tools.agent_types import AGENT_TYPE_MAPPING, AgentAudio, AgentImage, AgentText


if is_torch_available():
		import torch

if is_vision_available():
		from PIL import Image


_SCREAMING_SNAKE_CASE	       =      ["""text""", """image""", """audio"""]
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE )  ->     List[Any]:







							"""simple docstring"""


							__snake_case     			=       []

							for input_type in input_types:
														if input_type == "text":
																					inputs.append("Text input" )
														elif input_type == "image":
																					inputs.append(
																					    Image.open(Path(get_tests_dir("fixtures/tests_samples/COCO" ) ) / "000000039769.png" ).resize((5_12, 5_12) ) )
														elif input_type == "audio":
																					inputs.append(torch.ones(30_00 ) )
														elif isinstance(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE ):
																					inputs.append(create_inputs(SCREAMING_SNAKE_CASE ) )
														else:
																					raise ValueError(F'''Invalid type requested: {input_type}''' )

							return inputs
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE )  ->     Optional[Any]:







							"""simple docstring"""


							__snake_case     			=       []

							for output in outputs:
														if isinstance(SCREAMING_SNAKE_CASE							,    (str, AgentText) ):
																					output_types.append("text" )
														elif isinstance(SCREAMING_SNAKE_CASE							,    (Image.Image, AgentImage) ):
																					output_types.append("image" )
														elif isinstance(SCREAMING_SNAKE_CASE							,    (torch.Tensor, AgentAudio) ):
																					output_types.append("audio" )
														else:
																					raise ValueError(F'''Invalid output: {output}''' )

							return output_types







@is_tool_test
class 							__magic_name__     :





					def 							lowerCAmelCase      (				self   :							Optional[int]    ):
												self.assertTrue(hasattr(self.tool							,						"inputs"    )    )
												self.assertTrue(hasattr(self.tool							,						"outputs"    )    )

												__snake_case     			=       self.tool.inputs
												for _input in inputs:
																			if isinstance(_input							,						snake_case_    ):
																										for __input in _input:
																																	self.assertTrue(__input in authorized_types    )
																			else:
																										self.assertTrue(_input in authorized_types    )

												__snake_case     			=       self.tool.outputs
												for _output in outputs:
																			self.assertTrue(_output in authorized_types    )





					def 							lowerCAmelCase      (				self   :							int    ):
												__snake_case     			=       create_inputs(self.tool.inputs    )
												__snake_case     			=       self.tool(*snake_case_    )

												# There is a single output
												if len(self.tool.outputs    ) == 1:
																			__snake_case     			=       [outputs]

												self.assertListEqual(output_types(snake_case_    )							,						self.tool.outputs    )





					def 							lowerCAmelCase      (				self   :							Union[str, Any]    ):
												self.assertTrue(hasattr(self.tool							,						"description"    )    )
												self.assertTrue(hasattr(self.tool							,						"default_checkpoint"    )    )
												self.assertTrue(self.tool.description.startswith("This is a tool that"    )    )





					def 							lowerCAmelCase      (				self   :							Any    ):
												__snake_case     			=       create_inputs(self.tool.inputs    )
												__snake_case     			=       self.tool(*snake_case_    )

												if not isinstance(snake_case_							,						snake_case_    ):
																			__snake_case     			=       [outputs]

												self.assertEqual(len(snake_case_    )							,						len(self.tool.outputs    )    )

												for output, output_type in zip(snake_case_							,						self.tool.outputs    ):
																			__snake_case     			=       AGENT_TYPE_MAPPING[output_type]
																			self.assertTrue(isinstance(snake_case_							,						snake_case_    )    )





					def 							lowerCAmelCase      (				self   :							Tuple    ):
												__snake_case     			=       create_inputs(self.tool.inputs    )

												__snake_case     			=       []

												for _input, input_type in zip(snake_case_							,						self.tool.inputs    ):
																			if isinstance(snake_case_							,						snake_case_    ):
																										_inputs.append([AGENT_TYPE_MAPPING[_input_type](_input    ) for _input_type in input_type]    )
																			else:
																										_inputs.append(AGENT_TYPE_MAPPING[input_type](_input    )    )

        # Should not raise an error
												__snake_case     			=       self.tool(*snake_case_    )

												if not isinstance(snake_case_							,						snake_case_    ):
																			__snake_case     			=       [outputs]

												self.assertEqual(len(snake_case_    )							,						len(self.tool.outputs    )    )
