



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 			__snake_case   ( SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							UpperCamelCase__					:   PyTorchBenchmarkArguments
							UpperCamelCase__					:   PretrainedConfig
							UpperCamelCase__					:   str       =       "PyTorch"


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


							def _a   (				self	,	a_	,	a_	,	a_		):
											a__					   =		self._prepare_inference_func(a_	,	a_	,	a_		)
											return self._measure_speed(_inference		)


							def _a   (				self	,	a_	,	a_	,	a_		):
											a__					   =		self._prepare_inference_func(a_	,	a_	,	a_		)
											return self._measure_memory(_inference		)


							def _a   (				self	,	a_	,	a_	,	a_		):
											a__					   =		self._prepare_train_func(a_	,	a_	,	a_		)
											return self._measure_speed(_train		)


							def _a   (				self	,	a_	,	a_	,	a_		):
											a__					   =		self._prepare_train_func(a_	,	a_	,	a_		)
											return self._measure_memory(_train		)


							def _a   (				self	,	a_	,	a_	,	a_		):
											a__					   =		self.config_dict[model_name]

											if self.args.torchscript:
															a__					   =		True

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

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

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

											def encoder_decoder_forward():
															with torch.no_grad():
																			a__					   =		inference_model(a_	,	decoder_input_ids=a_		)
															return outputs

											def encoder_forward():
															with torch.no_grad():
																			a__					   =		inference_model(a_		)
															return outputs

											a__					   =		encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward


							def _a   (				self	,	a_	,	a_	,	a_		):
											a__					   =		self.config_dict[model_name]

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

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

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

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

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

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


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







							def _a   (				self	,	a_		):
											try:
															if self.args.trace_memory_line_by_line:
																			a__					   =		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."""		)
																							a__					   =		"""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()
																							a__					   =		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
																							a__					   =		nvml.nvmlDeviceGetMemoryInfo(a_		)
																							a__					   =		meminfo.used
																							a__					   =		Memory(a_		)
																							# shutdown nvml
																							nvml.nvmlShutdown()
															else:
																			# cpu
																			a__					   =		measure_peak_memory_cpu(a_		)
																			a__					   =		Memory(a_		) if isinstance(a_	,	a_		) else memory_bytes

															if self.args.trace_memory_line_by_line:
																			a__					   =		stop_memory_tracing(a_		)
															else:
																			a__					   =		None

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









import argparse

import torch

from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
from transformers.utils import logging


logging.set_verbosity_info()






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase     )     ->		str:







							"""simple docstring"""





							_a  :    Union[str, Any]			=						ElectraConfig.from_json_file(UpperCAmelCase     )
							print(F'Building PyTorch model from configuration: {config}'     )

							if discriminator_or_generator == "discriminator":
														_a  :    int			=						ElectraForPreTraining(UpperCAmelCase     )
							elif discriminator_or_generator == "generator":
														_a  :    Any			=						ElectraForMaskedLM(UpperCAmelCase     )
							else:
														raise ValueError('''The discriminator_or_generator argument should be either \'discriminator\' or \'generator\''''     )

							# Load weights from tf checkpoint
							load_tf_weights_in_electra(
							    UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase    ,	discriminator_or_generator=UpperCAmelCase     )

							# Save pytorch-model
							print(F'Save PyTorch model to {pytorch_dump_path}'     )
							torch.save(model.state_dict()    ,	UpperCAmelCase     )


if __name__ == "__main__":
						__lowerCamelCase						= argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
						)
						parser.add_argument(
						    '--config_file',
						    default=None,
						    type=str,
						    required=True,
						    help='The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.',
						)
						parser.add_argument(
						    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
						)
						parser.add_argument(
						    '--discriminator_or_generator',
						    default=None,
						    type=str,
						    required=True,
						    help=(
						        'Whether to export the generator or the discriminator. Should be a string, either \'discriminator\' or '
						        '\'generator\'.'
						    ),
						)
						__lowerCamelCase						= parser.parse_args()
						convert_tf_checkpoint_to_pytorch(
						    args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
						)