



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






					'''simple docstring'''





					_A						:     PyTorchBenchmarkArguments
					_A						:     PretrainedConfig
					_A						:     str		=	"PyTorch"






					@property
					def   lowerCAmelCase_							(			self:				str		)  ->      Optional[Any]:
											return torch.__version__






					def   lowerCAmelCase_							(			self:				Optional[int]		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      float:
											snake_case_					:Dict		=     self._prepare_inference_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_speed(_inference		)






					def   lowerCAmelCase_							(			self:				Tuple		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      [Memory, Optional[MemorySummary]]:
											snake_case_					:str		=     self._prepare_inference_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_memory(_inference		)






					def   lowerCAmelCase_							(			self:				Any		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      float:
											snake_case_					:int		=     self._prepare_train_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_speed(_train		)






					def   lowerCAmelCase_							(			self:				Any		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      [Memory, Optional[MemorySummary]]:
											snake_case_					:List[str]		=     self._prepare_train_func(snake_case		,	snake_case		,	snake_case		)
											return self._measure_memory(_train		)






					def   lowerCAmelCase_							(			self:				List[Any]		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      Callable[[], None]:
											snake_case_					:Union[str, Any]		=     self.config_dict[model_name]

											if self.args.torchscript:
																	snake_case_					:Any		=     True

											snake_case_					:List[Any]		=     (
											    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_					:Optional[int]		=     config.architectures[0]
																							snake_case_					:Union[str, Any]		=     __import__("""transformers"""		,	fromlist=[model_class]		)
																							snake_case_					:Dict		=     getattr(snake_case		,	snake_case		)
																							snake_case_					:Union[str, Any]		=     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_					:Any		=     MODEL_MAPPING[config.__class__](snake_case		)

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

											# encoder-decoder has vocab size saved differently
											snake_case_					:List[str]		=     config.vocab_size if hasattr(snake_case		,	"""vocab_size"""		) else config.encoder.vocab_size
											snake_case_					:Optional[Any]		=     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():
																							snake_case_					:Optional[Any]		=     torch.jit.trace(snake_case		,	snake_case		)
											else:
																	snake_case_					:Any		=     model

											def encoder_decoder_forward():
																	with torch.no_grad():
																							snake_case_					:List[str]		=     inference_model(snake_case		,	decoder_input_ids=snake_case		)
																	return outputs

											def encoder_forward():
																	with torch.no_grad():
																							snake_case_					:Optional[int]		=     inference_model(snake_case		)
																	return outputs

											snake_case_					:Dict		=     encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward






					def   lowerCAmelCase_							(			self:				Tuple		,	snake_case:				str		,	snake_case:				int		,	snake_case:				int		)  ->      Callable[[], None]:
											snake_case_					:Tuple		=     self.config_dict[model_name]

											snake_case_					:Dict		=     (
											    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_					:Optional[int]		=     config.architectures[0]
																							snake_case_					:List[str]		=     __import__("""transformers"""		,	fromlist=[model_class]		)
																							snake_case_					:Tuple		=     getattr(snake_case		,	snake_case		)
																							snake_case_					:str		=     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_					:Dict		=     MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case		)

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

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

											# encoder-decoder has vocab size saved differently
											snake_case_					:int		=     config.vocab_size if hasattr(snake_case		,	"""vocab_size"""		) else config.encoder.vocab_size
											snake_case_					:int		=     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():
																	snake_case_					:Optional[int]		=     train_model(snake_case		,	labels=snake_case		)[0]
																	loss.backward()
																	return loss

											def compute_loss_and_backprob_encoder_decoder():
																	snake_case_					:Any		=     train_model(snake_case		,	decoder_input_ids=snake_case		,	labels=snake_case		)[0]
																	loss.backward()
																	return loss

											snake_case_					:Optional[Any]		=     (
											    compute_loss_and_backprob_encoder_decoder
											    if config.is_encoder_decoder
											    else compute_loss_and_backprob_encoder
											)
											return _train






					def   lowerCAmelCase_							(			self:				List[Any]		,	snake_case:				Optional[Any]		)  ->      float:
											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
																	snake_case_					:Optional[int]		=     timeit.repeat(
																	    snake_case		,	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(snake_case		) / 1_0.0
											except RuntimeError as e:
																	self.print_fn(f"""Doesn't fit on GPU. {e}"""		)
																	return "N/A"






					def   lowerCAmelCase_							(			self:				Tuple		,	snake_case:				Callable[[], None]		)  ->      [Memory, MemorySummary]:
											try:
																	if self.args.trace_memory_line_by_line:
																							snake_case_					:Tuple		=     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."""		)
																													snake_case_					:Any		=     """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_					:Dict		=     nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
																													snake_case_					:str		=     nvml.nvmlDeviceGetMemoryInfo(snake_case		)
																													snake_case_					:Tuple		=     meminfo.used
																													snake_case_					:Optional[int]		=     Memory(snake_case		)
																													# shutdown nvml
																													nvml.nvmlShutdown()
																	else:
																							# cpu
																							snake_case_					:int		=     measure_peak_memory_cpu(snake_case		)
																							snake_case_					:List[str]		=     Memory(snake_case		) if isinstance(snake_case		,	snake_case		) else memory_bytes

																	if self.args.trace_memory_line_by_line:
																							snake_case_					:List[str]		=     stop_memory_tracing(snake_case		)
																	else:
																							snake_case_					:Tuple		=     None

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


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__UpperCAmelCase   						=     {
    "configuration_donut_swin": ["DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "DonutSwinConfig"],
    "processing_donut": ["DonutProcessor"],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __UpperCAmelCase   						=     [
       "DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
       "DonutSwinModel",
       "DonutSwinPreTrainedModel",
   ]

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __UpperCAmelCase   						=     ["DonutFeatureExtractor"]
   __UpperCAmelCase   						=     ["DonutImageProcessor"]


if TYPE_CHECKING:
   from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
   from .processing_donut import DonutProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_donut_swin import (
          DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
          DonutSwinModel,
          DonutSwinPreTrainedModel,
      )

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_donut import DonutFeatureExtractor
      from .image_processing_donut import DonutImageProcessor

else:
   import sys

   __UpperCAmelCase   						=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


