
"""simple docstring"""


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


A_						       =				logging.get_logger(__name__)




class    __lowerCamelCase     (   lowerCAmelCase ):
				a__:			int									=						42
				a__:			List[Any]									=						42
				a__:			Tuple									=						'PyTorch'




				@property
				def 					UpperCAmelCase__	(							self      ):
							return torch.__version__




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     self._prepare_inference_func(__lowerCamelCase       ,	__lowerCamelCase       ,	__lowerCamelCase      )
							return self._measure_speed(_inference      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     self._prepare_inference_func(__lowerCamelCase       ,	__lowerCamelCase       ,	__lowerCamelCase      )
							return self._measure_memory(_inference      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     self._prepare_train_func(__lowerCamelCase       ,	__lowerCamelCase       ,	__lowerCamelCase      )
							return self._measure_speed(_train      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     self._prepare_train_func(__lowerCamelCase       ,	__lowerCamelCase       ,	__lowerCamelCase      )
							return self._measure_memory(_train      )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     self.config_dict[model_name]

							if self.args.torchscript:
										lowerCamelCase_ 		=     True

							lowerCamelCase_ 		=     (
							    hasattr(__lowerCamelCase       ,	'''architectures'''      )
							    and isinstance(config.architectures       ,	__lowerCamelCase      )
							    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(__lowerCamelCase       ,	__lowerCamelCase      )
													lowerCamelCase_ 		=     model_cls(__lowerCamelCase      )
										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__](__lowerCamelCase      )

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

							# encoder-decoder has vocab size saved differently
							lowerCamelCase_ 		=     config.vocab_size if hasattr(__lowerCamelCase       ,	'''vocab_size'''      ) else config.encoder.vocab_size
							lowerCamelCase_ 		=     torch.randint(__lowerCamelCase       ,	(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(__lowerCamelCase       ,	__lowerCamelCase      )
							else:
										lowerCamelCase_ 		=     model

							def encoder_decoder_forward():
										with torch.no_grad():
													lowerCamelCase_ 		=     inference_model(__lowerCamelCase       ,	decoder_input_ids=__lowerCamelCase      )
										return outputs

							def encoder_forward():
										with torch.no_grad():
													lowerCamelCase_ 		=     inference_model(__lowerCamelCase      )
										return outputs

							lowerCamelCase_ 		=     encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
							return _forward




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     self.config_dict[model_name]

							lowerCamelCase_ 		=     (
							    hasattr(__lowerCamelCase       ,	'''architectures'''      )
							    and isinstance(config.architectures       ,	__lowerCamelCase      )
							    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(__lowerCamelCase       ,	__lowerCamelCase      )
													lowerCamelCase_ 		=     model_cls(__lowerCamelCase      )
										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__](__lowerCamelCase      )

							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(__lowerCamelCase       ,	'''vocab_size'''      ) else config.encoder.vocab_size
							lowerCamelCase_ 		=     torch.randint(__lowerCamelCase       ,	(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(__lowerCamelCase       ,	labels=__lowerCamelCase      )[0]
										loss.backward()
										return loss

							def compute_loss_and_backprob_encoder_decoder():
										lowerCamelCase_ 		=     train_model(__lowerCamelCase       ,	decoder_input_ids=__lowerCamelCase       ,	labels=__lowerCamelCase      )[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 					UpperCAmelCase__	(							self       ,	UpperCAmelCase      ):
							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(
													    __lowerCamelCase       ,	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(
										    __lowerCamelCase       ,	repeat=self.args.repeat       ,	number=10       ,	)

										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(__lowerCamelCase      ) / 10.0
							except RuntimeError as e:
										self.print_fn(f"Doesn't fit on GPU. {e}"      )
										return "N/A"





				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase      ):
							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(__lowerCamelCase      )
																lowerCamelCase_ 		=     meminfo.used
																lowerCamelCase_ 		=     Memory(__lowerCamelCase      )
																# shutdown nvml
																nvml.nvmlShutdown()
										else:
													# cpu
													lowerCamelCase_ 		=     measure_peak_memory_cpu(__lowerCamelCase      )
													lowerCamelCase_ 		=     Memory(__lowerCamelCase      ) if isinstance(__lowerCamelCase       ,	__lowerCamelCase      ) else memory_bytes

										if self.args.trace_memory_line_by_line:
													lowerCamelCase_ 		=     stop_memory_tracing(__lowerCamelCase      )
										else:
													lowerCamelCase_ 		=     None

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





"""simple docstring"""


import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import Mock, patch

import diffusers.utils.hub_utils




class    __lowerCamelCase     (   unittest.TestCase ):




				@patch('''diffusers.utils.hub_utils.get_full_repo_name'''      )
				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     '''full_repo_name'''
							with TemporaryDirectory() as tmpdir:
										# Dummy args values
										lowerCamelCase_ 		=     Mock()
										lowerCamelCase_ 		=     tmpdir
										lowerCamelCase_ 		=     0
										lowerCamelCase_ 		=     '''hub_token'''
										lowerCamelCase_ 		=     '''dataset_name'''
										lowerCamelCase_ 		=     0.0_1
										lowerCamelCase_ 		=     10_0000
										lowerCamelCase_ 		=     1_0000
										lowerCamelCase_ 		=     0.0_1
										lowerCamelCase_ 		=     0.0_2
										lowerCamelCase_ 		=     0.0_3
										lowerCamelCase_ 		=     0.0_0_0_5
										lowerCamelCase_ 		=     0.0_0_0_0_0_1
										lowerCamelCase_ 		=     1
										lowerCamelCase_ 		=     10
										lowerCamelCase_ 		=     0.0_0_1
										lowerCamelCase_ 		=     0.1
										lowerCamelCase_ 		=     0.2
										lowerCamelCase_ 		=     True

										# Model card mush be rendered and saved
										diffusers.utils.hub_utils.create_model_card(UpperCAmelCase       ,	model_name='''model_name'''      )
										self.assertTrue((Path(UpperCAmelCase      ) / '''README.md''').is_file()      )



