
"""simple docstring"""




from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


snake_case			       =		logging.get_logger(__name__)

snake_case			       =		{
    'apple/mobilevit-small': 'https://huggingface.co/apple/mobilevit-small/resolve/main/config.json',
    'apple/mobilevit-x-small': 'https://huggingface.co/apple/mobilevit-x-small/resolve/main/config.json',
    'apple/mobilevit-xx-small': 'https://huggingface.co/apple/mobilevit-xx-small/resolve/main/config.json',
    'apple/deeplabv3-mobilevit-small': (
        'https://huggingface.co/apple/deeplabv3-mobilevit-small/resolve/main/config.json'
    ),
    'apple/deeplabv3-mobilevit-x-small': (
        'https://huggingface.co/apple/deeplabv3-mobilevit-x-small/resolve/main/config.json'
    ),
    'apple/deeplabv3-mobilevit-xx-small': (
        'https://huggingface.co/apple/deeplabv3-mobilevit-xx-small/resolve/main/config.json'
    ),
    # See all MobileViT models at https://huggingface.co/models?filter=mobilevit
}




class 					UpperCamelCase     (	__magic_name__				):

				"""simple docstring"""



				UpperCAmelCase_    :				List[str]						  =    "mobilevit"


				def __init__(    self		,     lowercase__=3		,     lowercase__=256		,     lowercase__=2		,     lowercase__=[144, 192, 240]		,     lowercase__=[16, 32, 64, 96, 128, 160, 640]		,     lowercase__=4		,     lowercase__=2.0		,     lowercase__=4.0		,     lowercase__="silu"		,     lowercase__=3		,     lowercase__=32		,     lowercase__=0.1		,     lowercase__=0.0		,     lowercase__=0.1		,     lowercase__=0.02		,     lowercase__=1E-5		,     lowercase__=True		,     lowercase__=256		,     lowercase__=[6, 12, 18]		,     lowercase__=0.1		,     lowercase__=255		,     **lowercase__		,     )   ->   Dict:
									"""simple docstring"""







									super().__init__(**lowercase__      )

									SCREAMING_SNAKE_CASE           =     num_channels
									SCREAMING_SNAKE_CASE           =     image_size
									SCREAMING_SNAKE_CASE           =     patch_size
									SCREAMING_SNAKE_CASE           =     hidden_sizes
									SCREAMING_SNAKE_CASE           =     neck_hidden_sizes
									SCREAMING_SNAKE_CASE           =     num_attention_heads
									SCREAMING_SNAKE_CASE           =     mlp_ratio
									SCREAMING_SNAKE_CASE           =     expand_ratio
									SCREAMING_SNAKE_CASE           =     hidden_act
									SCREAMING_SNAKE_CASE           =     conv_kernel_size
									SCREAMING_SNAKE_CASE           =     output_stride
									SCREAMING_SNAKE_CASE           =     hidden_dropout_prob
									SCREAMING_SNAKE_CASE           =     attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE           =     classifier_dropout_prob
									SCREAMING_SNAKE_CASE           =     initializer_range
									SCREAMING_SNAKE_CASE           =     layer_norm_eps
									SCREAMING_SNAKE_CASE           =     qkv_bias

									# decode head attributes for semantic segmentation
									SCREAMING_SNAKE_CASE           =     aspp_out_channels
									SCREAMING_SNAKE_CASE           =     atrous_rates
									SCREAMING_SNAKE_CASE           =     aspp_dropout_prob
									SCREAMING_SNAKE_CASE           =     semantic_loss_ignore_index






class 					UpperCamelCase     (	__magic_name__				):

				"""simple docstring"""



				UpperCAmelCase_    :				Union[str, Any]						  =    version.parse("1.11"				)


				@property
				def 							A (    self      )   ->   Mapping[str, Mapping[int, str]]:
									"""simple docstring"""







									return OrderedDict([('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'})]      )


				@property
				def 							A (    self      )   ->   Mapping[str, Mapping[int, str]]:
									"""simple docstring"""







									if self.task == "image-classification":
														return OrderedDict([('logits', {0: 'batch'})]      )
									else:
														return OrderedDict([('last_hidden_state', {0: 'batch'}), ('pooler_output', {0: 'batch'})]      )




				@property
				def 							A (    self      )   ->   float:
									"""simple docstring"""







									return 1E-4


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


snake_case			       =		logging.get_logger(__name__)




class 					UpperCamelCase     (	__magic_name__				):

				"""simple docstring"""



				UpperCAmelCase_    :				PyTorchBenchmarkArguments
				UpperCAmelCase_    :				PretrainedConfig
				UpperCAmelCase_    :				str						  =    "PyTorch"


				@property
				def 							A (    self      )   ->   List[str]:
									"""simple docstring"""







									return torch.__version__


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   float:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self._prepare_inference_func(lowercase__		,     lowercase__		,     lowercase__      )
									return self._measure_speed(_inference      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   [Memory, Optional[MemorySummary]]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self._prepare_inference_func(lowercase__		,     lowercase__		,     lowercase__      )
									return self._measure_memory(_inference      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   float:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self._prepare_train_func(lowercase__		,     lowercase__		,     lowercase__      )
									return self._measure_speed(_train      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   [Memory, Optional[MemorySummary]]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self._prepare_train_func(lowercase__		,     lowercase__		,     lowercase__      )
									return self._measure_memory(_train      )


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   Callable[[], None]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.config_dict[model_name]

									if self.args.torchscript:
														SCREAMING_SNAKE_CASE           =     True

									SCREAMING_SNAKE_CASE           =     (
									    hasattr(lowercase__		,     'architectures'      )
									    and isinstance(config.architectures		,     lowercase__      )
									    and len(config.architectures      ) > 0
									)
									if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																			SCREAMING_SNAKE_CASE           =     config.architectures[0]
																			SCREAMING_SNAKE_CASE           =     __import__('transformers'		,     fromlist=[model_class]      )
																			SCREAMING_SNAKE_CASE           =     getattr(lowercase__		,     lowercase__      )
																			SCREAMING_SNAKE_CASE           =     model_cls(lowercase__      )
														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:
														SCREAMING_SNAKE_CASE           =     MODEL_MAPPING[config.__class__](lowercase__      )

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

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

									def encoder_decoder_forward():
														with torch.no_grad():
																			SCREAMING_SNAKE_CASE           =     inference_model(lowercase__		,     decoder_input_ids=lowercase__      )
														return outputs

									def encoder_forward():
														with torch.no_grad():
																			SCREAMING_SNAKE_CASE           =     inference_model(lowercase__      )
														return outputs

									SCREAMING_SNAKE_CASE           =     encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
									return _forward


				def 							A (    self		,     lowercase__		,     lowercase__		,     lowercase__      )   ->   Callable[[], None]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     self.config_dict[model_name]

									SCREAMING_SNAKE_CASE           =     (
									    hasattr(lowercase__		,     'architectures'      )
									    and isinstance(config.architectures		,     lowercase__      )
									    and len(config.architectures      ) > 0
									)
									if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																			SCREAMING_SNAKE_CASE           =     config.architectures[0]
																			SCREAMING_SNAKE_CASE           =     __import__('transformers'		,     fromlist=[model_class]      )
																			SCREAMING_SNAKE_CASE           =     getattr(lowercase__		,     lowercase__      )
																			SCREAMING_SNAKE_CASE           =     model_cls(lowercase__      )
														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:
														SCREAMING_SNAKE_CASE           =     MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowercase__      )

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

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

									# encoder-decoder has vocab size saved differently
									SCREAMING_SNAKE_CASE           =     config.vocab_size if hasattr(lowercase__		,     'vocab_size'      ) else config.encoder.vocab_size
									SCREAMING_SNAKE_CASE           =     torch.randint(lowercase__		,     (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():
														SCREAMING_SNAKE_CASE           =     train_model(lowercase__		,     labels=lowercase__      )[0]
														loss.backward()
														return loss

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

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


				def 							A (    self		,     lowercase__      )   ->   float:
									"""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(
																			    lowercase__		,     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
														SCREAMING_SNAKE_CASE           =     timeit.repeat(
														    lowercase__		,     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(lowercase__      ) / 10.0
									except RuntimeError as e:
														self.print_fn(f'''Doesn\'t fit on GPU. {e}'''      )
														return "N/A"




				def 							A (    self		,     lowercase__      )   ->   [Memory, MemorySummary]:
									"""simple docstring"""







									try:
														if self.args.trace_memory_line_by_line:
																			SCREAMING_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'
																			    ' `--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.'      )
																								SCREAMING_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()
																								SCREAMING_SNAKE_CASE           =     nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx      )
																								SCREAMING_SNAKE_CASE           =     nvml.nvmlDeviceGetMemoryInfo(lowercase__      )
																								SCREAMING_SNAKE_CASE           =     meminfo.used
																								SCREAMING_SNAKE_CASE           =     Memory(lowercase__      )
																								# shutdown nvml
																								nvml.nvmlShutdown()
														else:
																			# cpu
																			SCREAMING_SNAKE_CASE           =     measure_peak_memory_cpu(lowercase__      )
																			SCREAMING_SNAKE_CASE           =     Memory(lowercase__      ) if isinstance(lowercase__		,     lowercase__      ) else memory_bytes

														if self.args.trace_memory_line_by_line:
																			SCREAMING_SNAKE_CASE           =     stop_memory_tracing(lowercase__      )
														else:
																			SCREAMING_SNAKE_CASE           =     None

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