




import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
					import torch

					from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
					import pyanvml.pyanvml as nvml


UpperCamelCase__            =     logging.get_logger(__name__)
class      a__       (   snake_case__ ):
						_a				:       PyTorchBenchmarkArguments
						_a				:       PretrainedConfig
						_a				:       str   =       "PyTorch"






						@property
						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											return torch.__version__






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			self._prepare_inference_func(_snake_case						,     _snake_case						,     _snake_case     )
											return self._measure_speed(_inference     )






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			self._prepare_inference_func(_snake_case						,     _snake_case						,     _snake_case     )
											return self._measure_memory(_inference     )






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			self._prepare_train_func(_snake_case						,     _snake_case						,     _snake_case     )
											return self._measure_speed(_train     )






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			self._prepare_train_func(_snake_case						,     _snake_case						,     _snake_case     )
											return self._measure_memory(_train     )






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			self.config_dict[model_name]

											if self.args.torchscript:
																__lowerCAmelCase							  =			True

											__lowerCAmelCase							  =			(
											    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:
																					__lowerCAmelCase							  =			config.architectures[0]
																					__lowerCAmelCase							  =			__import__("transformers"						,     fromlist=[model_class]     )
																					__lowerCAmelCase							  =			getattr(_snake_case						,     _snake_case     )
																					__lowerCAmelCase							  =			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:
																__lowerCAmelCase							  =			MODEL_MAPPING[config.__class__](_snake_case     )

											model.eval()
											model.to(self.args.device     )

											# encoder-decoder has vocab size saved differently
											__lowerCAmelCase							  =			config.vocab_size if hasattr(_snake_case						,     "vocab_size"     ) else config.encoder.vocab_size
											__lowerCAmelCase							  =			torch.randint(_snake_case						,     (batch_size, sequence_length)						,     dtype=torch.long						,     device=self.args.device     )

											if self.args.fpaa:
																logger.info("Running training in Mixed Precision..."     )
																if not self.args.is_gpu:
																					raise ValueError("Mixed precision is possible only for GPU."     )
																# amp seems to have memory leaks so that memory usage
																# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
																model.half()

											if self.args.torchscript:
																with torch.no_grad():
																					__lowerCAmelCase							  =			torch.jit.trace(_snake_case						,     _snake_case     )
											else:
																__lowerCAmelCase							  =			model

											def encoder_decoder_forward():
																with torch.no_grad():
																					__lowerCAmelCase							  =			inference_model(_snake_case						,     decoder_input_ids=_snake_case     )
																return outputs

											def encoder_forward():
																with torch.no_grad():
																					__lowerCAmelCase							  =			inference_model(_snake_case     )
																return outputs

											__lowerCAmelCase							  =			encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			self.config_dict[model_name]

											__lowerCAmelCase							  =			(
											    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:
																					__lowerCAmelCase							  =			config.architectures[0]
																					__lowerCAmelCase							  =			__import__("transformers"						,     fromlist=[model_class]     )
																					__lowerCAmelCase							  =			getattr(_snake_case						,     _snake_case     )
																					__lowerCAmelCase							  =			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:
																__lowerCAmelCase							  =			MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_snake_case     )

											if self.args.torchscript:
																raise NotImplementedError("Training for torchscript is currently not implemented"     )
											else:
																__lowerCAmelCase							  =			model

											model.train()
											model.to(self.args.device     )

											# encoder-decoder has vocab size saved differently
											__lowerCAmelCase							  =			config.vocab_size if hasattr(_snake_case						,     "vocab_size"     ) else config.encoder.vocab_size
											__lowerCAmelCase							  =			torch.randint(_snake_case						,     (batch_size, sequence_length)						,     dtype=torch.long						,     device=self.args.device     )

											if self.args.fpaa:
																logger.info("Running training in Mixed Precision..."     )
																if not self.args.is_gpu:
																					raise ValueError("Mixed precision is possible only for GPU."     )

																# amp seems to have memory leaks so that memory usage
																# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
																model.half()

											def compute_loss_and_backprob_encoder():
																__lowerCAmelCase							  =			train_model(_snake_case						,     labels=_snake_case     )[0]
																loss.backward()
																return loss

											def compute_loss_and_backprob_encoder_decoder():
																__lowerCAmelCase							  =			train_model(_snake_case						,     decoder_input_ids=_snake_case						,     labels=_snake_case     )[0]
																loss.backward()
																return loss

											__lowerCAmelCase							  =			(
											    compute_loss_and_backprob_encoder_decoder
											    if config.is_encoder_decoder
											    else compute_loss_and_backprob_encoder
											)
											return _train






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A     ):




											"""simple docstring"""



											try:
																if self.args.is_tpu or self.args.torchscript:
																					# run additional 10 times to stabilize compilation for tpu and torchscript
																					logger.info("Do inference on TPU or torchscript. 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
																__lowerCAmelCase							  =			timeit.repeat(
																    _snake_case						,     repeat=self.args.repeat						,     number=1_0						,     )

																if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
																					import torch_xla.debug.metrics as met

																					self.print_fn(met.metrics_report()     )

																return min(_snake_case     ) / 10.0
											except RuntimeError as e:
																self.print_fn(f"""Doesn\'t fit on GPU. {e}"""     )
																return "N/A"






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A     ):




											"""simple docstring"""



											try:
																if self.args.trace_memory_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"
																					    " `--no-memory` or `args.memory=False`"     )
																elif self.args.is_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(_snake_case     )
																										__lowerCAmelCase							  =			meminfo.used
																										__lowerCAmelCase							  =			Memory(_snake_case     )
																										# shutdown nvml
																										nvml.nvmlShutdown()
																else:
																					# cpu
																					__lowerCAmelCase							  =			measure_peak_memory_cpu(_snake_case     )
																					__lowerCAmelCase							  =			Memory(_snake_case     ) if isinstance(_snake_case						,     _snake_case     ) else memory_bytes

																if self.args.trace_memory_line_by_line:
																					__lowerCAmelCase							  =			stop_memory_tracing(_snake_case     )
																else:
																					__lowerCAmelCase							  =			None

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












import unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types




class 	lowerCamelCase__				(							unittest.TestCase,   UpperCAmelCase ):






      def      UpperCAmelCase_     (self					:    Any		)       ->   Union[str, Any]:


       """simple docstring"""






       lowerCamelCase_							:    str							=     load_tool('translation'		)
       self.tool.setup()
       lowerCamelCase_							:    Optional[int]							=     load_tool('translation'		,		remote=_snake_case		)






      def      UpperCAmelCase_     (self					:    Any		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_							:    Optional[Any]							=     self.tool('Hey, what\'s up?'		,		src_lang='English'		,		tgt_lang='French'		)
       self.assertEqual(_snake_case		,		'- Hé, comment ça va?'		)






      def      UpperCAmelCase_     (self					:    Optional[int]		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_							:    Optional[int]							=     self.remote_tool('Hey, what\'s up?'		,		src_lang='English'		,		tgt_lang='French'		)
       self.assertEqual(_snake_case		,		'- Hé, comment ça va?'		)






      def      UpperCAmelCase_     (self					:    List[str]		)       ->   Optional[int]:


       """simple docstring"""






       lowerCamelCase_							:    Union[str, Any]							=     self.tool(text='Hey, what\'s up?'		,		src_lang='English'		,		tgt_lang='French'		)
       self.assertEqual(_snake_case		,		'- Hé, comment ça va?'		)






      def      UpperCAmelCase_     (self					:    Any		)       ->   Tuple:


       """simple docstring"""






       lowerCamelCase_							:    Optional[int]							=     self.remote_tool(text='Hey, what\'s up?'		,		src_lang='English'		,		tgt_lang='French'		)
       self.assertEqual(_snake_case		,		'- Hé, comment ça va?'		)






      def      UpperCAmelCase_     (self					:    Union[str, Any]		)       ->   Any:


       """simple docstring"""






       lowerCamelCase_							:    str							=     ['Hey, what\'s up?', 'English', 'Spanish']
       lowerCamelCase_							:    Any							=     self.tool(*_snake_case		)

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

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






      def      UpperCAmelCase_     (self					:    Tuple		)       ->   Optional[Any]:


       """simple docstring"""






       lowerCamelCase_							:    List[str]							=     ['Hey, what\'s up?', 'English', 'Spanish']
       lowerCamelCase_							:    Optional[Any]							=     self.tool(*_snake_case		)

       if not isinstance(_snake_case		,		_snake_case		):
        lowerCamelCase_							:    List[str]							=     [outputs]

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

       for output, output_type in zip(_snake_case		,		self.tool.outputs		):
        lowerCamelCase_							:    Optional[int]							=     AGENT_TYPE_MAPPING[output_type]
        self.assertTrue(isinstance(_snake_case		,		_snake_case		)		)






      def      UpperCAmelCase_     (self					:    Any		)       ->   Optional[Any]:


       """simple docstring"""






       lowerCamelCase_							:    Dict							=     ['Hey, what\'s up?', 'English', 'Spanish']

       lowerCamelCase_							:    List[str]							=     []

       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
       lowerCamelCase_							:    Any							=     self.tool(*_snake_case		)

       if not isinstance(_snake_case		,		_snake_case		):
        lowerCamelCase_							:    Optional[Any]							=     [outputs]

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