

'''simple docstring'''





import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


lowerCAmelCase__								=		logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm
							import torch_xla.debug.metrics as met







class 	lowercase    (				_snake_case							):


		"""simple docstring"""




		def __init__(				self	,       *__snake_case	,       __snake_case=None	,       __snake_case=None	,       __snake_case=None	,       **__snake_case):
			super().__init__(*lowerCAmelCase__	,       **lowerCAmelCase__)
			_UpperCamelCase   :		List[str]			=						eval_examples
			_UpperCamelCase   :		Dict			=						post_process_function
			_UpperCamelCase   :		Union[str, Any]			=						quant_trainer_args
			_UpperCamelCase   :		Any			=						1_28  # default number of calibration samples


		def 		A__					(				self	,       __snake_case=None):
			if calib_dataset is None and self.calib_dataset is None:
				raise ValueError('Trainer: calibration requires an calib_dataset.')
			_UpperCamelCase   :		Optional[int]			=						calib_dataset if calib_dataset is not None else self.calib_dataset

			_UpperCamelCase   :		str			=						self._remove_unused_columns(lowerCAmelCase__	,       description='Calibration')

			return DataLoader(
			    lowerCAmelCase__	,       batch_size=self.args.eval_batch_size	,       collate_fn=self.data_collator	,       drop_last=self.args.dataloader_drop_last	,       num_workers=self.args.dataloader_num_workers	,       pin_memory=self.args.dataloader_pin_memory	,       shuffle=lowerCAmelCase__	,       )


		def 		A__					(				self	,       __snake_case=None):
			_UpperCamelCase   :		Optional[int]			=						self.train_dataset if calib_dataset is None else calib_dataset
			_UpperCamelCase   :		Optional[int]			=						self.get_calib_dataloader(lowerCAmelCase__)

			_UpperCamelCase   :		Tuple			=						self.model
			quant_trainer.configure_model(lowerCAmelCase__	,       self.quant_trainer_args	,       calib=lowerCAmelCase__)
			model.eval()
			quant_trainer.enable_calibration(lowerCAmelCase__)

			logger.info('***** Running calibration *****')
			logger.info(f'''  Num examples = {self.calib_num}''')
			logger.info(f'''  Batch size = {calib_dataloader.batch_size}''')

			for step, inputs in enumerate(lowerCAmelCase__):
				# Prediction step
				_UpperCamelCase	,		_UpperCamelCase	,		_UpperCamelCase   :		List[str]			=						self.prediction_step(lowerCAmelCase__	,       lowerCAmelCase__	,       prediction_loss_only=lowerCAmelCase__)
				if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
					break

			quant_trainer.finish_calibration(lowerCAmelCase__	,       self.quant_trainer_args)
			_UpperCamelCase   :		List[Any]			=						model


		def 		A__					(				self	,       __snake_case=None	,       __snake_case=None	,       __snake_case=None	,       __snake_case = "eval"):
			_UpperCamelCase   :		str			=						self.eval_dataset if eval_dataset is None else eval_dataset
			_UpperCamelCase   :		str			=						self.get_eval_dataloader(lowerCAmelCase__)
			_UpperCamelCase   :		Dict			=						self.eval_examples if eval_examples is None else eval_examples

			# Temporarily disable metric computation, we will do it in the loop here.
			_UpperCamelCase   :		Any			=						self.compute_metrics
			_UpperCamelCase   :		List[str]			=						None
			_UpperCamelCase   :		Dict			=						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
			try:
				_UpperCamelCase   :		List[Any]			=						eval_loop(
				    lowerCAmelCase__	,       description='Evaluation'	,       prediction_loss_only=True if compute_metrics is None else None	,       ignore_keys=lowerCAmelCase__	,       )
			finally:
				_UpperCamelCase   :		Tuple			=						compute_metrics

			if self.post_process_function is not None and self.compute_metrics is not None:
				_UpperCamelCase   :		Tuple			=						self.post_process_function(lowerCAmelCase__	,       lowerCAmelCase__	,       output.predictions)
				_UpperCamelCase   :		Dict			=						self.compute_metrics(lowerCAmelCase__)

				# Prefix all keys with metric_key_prefix + '_'
				for key in list(metrics.keys()):
					if not key.startswith(f'''{metric_key_prefix}_'''):
						_UpperCamelCase   :		Optional[Any]			=						metrics.pop(lowerCAmelCase__)

				self.log(lowerCAmelCase__)
			else:
				_UpperCamelCase   :		Tuple			=						{}

			if self.args.tpu_metrics_debug or self.args.debug:
				# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
				xm.master_print(met.metrics_report())

			_UpperCamelCase   :		int			=						self.callback_handler.on_evaluate(self.args	,       self.state	,       self.control	,       lowerCAmelCase__)
			return metrics


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case=None	,       __snake_case = "test"):
			_UpperCamelCase   :		Optional[int]			=						self.get_test_dataloader(lowerCAmelCase__)

			# Temporarily disable metric computation, we will do it in the loop here.
			_UpperCamelCase   :		str			=						self.compute_metrics
			_UpperCamelCase   :		str			=						None
			_UpperCamelCase   :		List[str]			=						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
			try:
				_UpperCamelCase   :		int			=						eval_loop(
				    lowerCAmelCase__	,       description='Prediction'	,       prediction_loss_only=True if compute_metrics is None else None	,       ignore_keys=lowerCAmelCase__	,       )
			finally:
				_UpperCamelCase   :		Union[str, Any]			=						compute_metrics

			if self.post_process_function is None or self.compute_metrics is None:
				return output

			_UpperCamelCase   :		Union[str, Any]			=						self.post_process_function(lowerCAmelCase__	,       lowerCAmelCase__	,       output.predictions	,       'predict')
			_UpperCamelCase   :		Optional[Any]			=						self.compute_metrics(lowerCAmelCase__)

			# Prefix all keys with metric_key_prefix + '_'
			for key in list(metrics.keys()):
				if not key.startswith(f'''{metric_key_prefix}_'''):
					_UpperCamelCase   :		str			=						metrics.pop(lowerCAmelCase__)

			return PredictionOutput(predictions=predictions.predictions	,       label_ids=predictions.label_ids	,       metrics=lowerCAmelCase__)






		def 		A__					(				self	,       __snake_case="./"):
			_UpperCamelCase   :		List[Any]			=						self.eval_dataset
			_UpperCamelCase   :		Optional[int]			=						self.get_eval_dataloader(lowerCAmelCase__)

			_UpperCamelCase   :		str			=						next(iter(lowerCAmelCase__))

			# saving device - to make it consistent
			_UpperCamelCase   :		Dict			=						torch.device('cuda' if torch.cuda.is_available() else 'cpu')

			# convert to tuple
			_UpperCamelCase   :		Optional[int]			=						tuple(v.to(lowerCAmelCase__) for k, v in batch.items())

			logger.info('Converting model to be onnx compatible')
			from pytorch_quantization.nn import TensorQuantizer

			_UpperCamelCase   :		int			=						True

			_UpperCamelCase   :		Any			=						self.model.to(lowerCAmelCase__)

			model.eval()
			model.float()

			_UpperCamelCase   :		str			=						model.module if hasattr(lowerCAmelCase__	,       'module') else model
			quant_trainer.configure_model(lowerCAmelCase__	,       self.quant_trainer_args)

			_UpperCamelCase   :		Optional[int]			=						os.path.join(lowerCAmelCase__	,       'model.onnx')
			logger.info(f'''exporting model to {output_model_file}''')

			_UpperCamelCase   :		Optional[int]			=						{0: 'batch_size', 1: 'seq_len'}

			torch.onnx.export(
			    lowerCAmelCase__	,       lowerCAmelCase__	,       lowerCAmelCase__	,       export_params=lowerCAmelCase__	,       opset_version=13	,       do_constant_folding=lowerCAmelCase__	,       input_names=['input_ids', 'attention_mask', 'token_type_ids']	,       output_names=['output_start_logits', 'output_end_logits']	,       dynamic_axes={
			        'input_ids': axes,
			        'attention_mask': axes,
			        'token_type_ids': axes,
			        'output_start_logits': axes,
			        'output_end_logits': axes,
			    }	,       verbose=lowerCAmelCase__	,       )
			logger.info('onnx export finished')







import logging

from transformers.configuration_utils import PretrainedConfig


lowerCAmelCase__								=		logging.getLogger(__name__)







class 	lowercase    (				_lowercase							):


		"""simple docstring"""




		a__					       =  "masked_bert"


		def __init__(				self	,       __snake_case=3_05_22	,       __snake_case=7_68	,       __snake_case=12	,       __snake_case=12	,       __snake_case=30_72	,       __snake_case="gelu"	,       __snake_case=0.1	,       __snake_case=0.1	,       __snake_case=5_12	,       __snake_case=2	,       __snake_case=0.0_2	,       __snake_case=1e-12	,       __snake_case=0	,       __snake_case="topK"	,       __snake_case="constant"	,       __snake_case=0.0	,       **__snake_case	,       ):
			super().__init__(pad_token_id=__snake_case	,       **__snake_case)

			_UpperCamelCase   :		List[Any]			=						vocab_size
			_UpperCamelCase   :		Union[str, Any]			=						hidden_size
			_UpperCamelCase   :		Optional[int]			=						num_hidden_layers
			_UpperCamelCase   :		Any			=						num_attention_heads
			_UpperCamelCase   :		int			=						hidden_act
			_UpperCamelCase   :		str			=						intermediate_size
			_UpperCamelCase   :		str			=						hidden_dropout_prob
			_UpperCamelCase   :		Any			=						attention_probs_dropout_prob
			_UpperCamelCase   :		Tuple			=						max_position_embeddings
			_UpperCamelCase   :		Dict			=						type_vocab_size
			_UpperCamelCase   :		str			=						initializer_range
			_UpperCamelCase   :		List[Any]			=						layer_norm_eps
			_UpperCamelCase   :		Tuple			=						pruning_method
			_UpperCamelCase   :		Tuple			=						mask_init
			_UpperCamelCase   :		Dict			=						mask_scale




