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


__lowercase        =    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 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			def __init__(  self			,       *__lowerCAmelCase			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       **__lowerCAmelCase):




					"""simple docstring"""


					super().__init__(*__lowerCAmelCase			,       **__lowerCAmelCase)
					lowerCAmelCase     =    eval_examples
					lowerCAmelCase     =    post_process_function
					lowerCAmelCase     =    quant_trainer_args
					lowerCAmelCase     =    128  # default number of calibration samples



			def 				a_			(  self			,       __lowerCAmelCase=None):




					"""simple docstring"""


					if calib_dataset is None and self.calib_dataset is None:
							raise ValueError("""Trainer: calibration requires an calib_dataset.""")
					lowerCAmelCase     =    calib_dataset if calib_dataset is not None else self.calib_dataset

					lowerCAmelCase     =    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			,       __lowerCAmelCase=None):




					"""simple docstring"""


					lowerCAmelCase     =    self.train_dataset if calib_dataset is None else calib_dataset
					lowerCAmelCase     =    self.get_calib_dataloader(__lowerCAmelCase)

					lowerCAmelCase     =    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
							lowerCAmelCase   ,	lowerCAmelCase   ,	lowerCAmelCase     =    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)
					lowerCAmelCase     =    model



			def 				a_			(  self			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       __lowerCAmelCase = "eval"):




					"""simple docstring"""


					lowerCAmelCase     =    self.eval_dataset if eval_dataset is None else eval_dataset
					lowerCAmelCase     =    self.get_eval_dataloader(__lowerCAmelCase)
					lowerCAmelCase     =    self.eval_examples if eval_examples is None else eval_examples

					# Temporarily disable metric computation, we will do it in the loop here.
					lowerCAmelCase     =    self.compute_metrics
					lowerCAmelCase     =    None
					lowerCAmelCase     =    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
							lowerCAmelCase     =    eval_loop(
							    __lowerCAmelCase			,       description="""Evaluation"""			,       prediction_loss_only=True if compute_metrics is None else None			,       ignore_keys=__lowerCAmelCase			,       )
					finally:
							lowerCAmelCase     =    compute_metrics

					if self.post_process_function is not None and self.compute_metrics is not None:
							lowerCAmelCase     =    self.post_process_function(__lowerCAmelCase			,       __lowerCAmelCase			,       output.predictions)
							lowerCAmelCase     =    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}_"):
											lowerCAmelCase     =    metrics.pop(__lowerCAmelCase)

							self.log(__lowerCAmelCase)
					else:
							lowerCAmelCase     =    {}

					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())

					lowerCAmelCase     =    self.callback_handler.on_evaluate(self.args			,       self.state			,       self.control			,       __lowerCAmelCase)
					return metrics



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase=None			,       __lowerCAmelCase = "test"):




					"""simple docstring"""


					lowerCAmelCase     =    self.get_test_dataloader(__lowerCAmelCase)

					# Temporarily disable metric computation, we will do it in the loop here.
					lowerCAmelCase     =    self.compute_metrics
					lowerCAmelCase     =    None
					lowerCAmelCase     =    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
							lowerCAmelCase     =    eval_loop(
							    __lowerCAmelCase			,       description="""Prediction"""			,       prediction_loss_only=True if compute_metrics is None else None			,       ignore_keys=__lowerCAmelCase			,       )
					finally:
							lowerCAmelCase     =    compute_metrics

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

					lowerCAmelCase     =    self.post_process_function(__lowerCAmelCase			,       __lowerCAmelCase			,       output.predictions			,       """predict""")
					lowerCAmelCase     =    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}_"):
									lowerCAmelCase     =    metrics.pop(__lowerCAmelCase)

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







			def 				a_			(  self			,       __lowerCAmelCase="./"):




					"""simple docstring"""


					lowerCAmelCase     =    self.eval_dataset
					lowerCAmelCase     =    self.get_eval_dataloader(__lowerCAmelCase)

					lowerCAmelCase     =    next(iter(__lowerCAmelCase))

					# saving device - to make it consistent
					lowerCAmelCase     =    torch.device("""cuda""" if torch.cuda.is_available() else """cpu""")

					# convert to tuple
					lowerCAmelCase     =    tuple(v.to(__lowerCAmelCase) for k, v in batch.items())

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

					lowerCAmelCase     =    True

					lowerCAmelCase     =    self.model.to(__lowerCAmelCase)

					model.eval()
					model.float()

					lowerCAmelCase     =    model.module if hasattr(__lowerCAmelCase			,       """module""") else model
					quant_trainer.configure_model(__lowerCAmelCase			,       self.quant_trainer_args)

					lowerCAmelCase     =    os.path.join(__lowerCAmelCase			,       """model.onnx""")
					logger.info(f"exporting model to {output_model_file}")

					lowerCAmelCase     =    {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""")





'''simple docstring'''




from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig


__lowercase        =    {
    '''albert-base-v1''': '''https://huggingface.co/albert-base-v1/resolve/main/config.json''',
    '''albert-large-v1''': '''https://huggingface.co/albert-large-v1/resolve/main/config.json''',
    '''albert-xlarge-v1''': '''https://huggingface.co/albert-xlarge-v1/resolve/main/config.json''',
    '''albert-xxlarge-v1''': '''https://huggingface.co/albert-xxlarge-v1/resolve/main/config.json''',
    '''albert-base-v2''': '''https://huggingface.co/albert-base-v2/resolve/main/config.json''',
    '''albert-large-v2''': '''https://huggingface.co/albert-large-v2/resolve/main/config.json''',
    '''albert-xlarge-v2''': '''https://huggingface.co/albert-xlarge-v2/resolve/main/config.json''',
    '''albert-xxlarge-v2''': '''https://huggingface.co/albert-xxlarge-v2/resolve/main/config.json''',
}






class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	List[str]						=			'''albert'''



			def __init__(  self			,       __lowerCAmelCase=30000			,       __lowerCAmelCase=128			,       __lowerCAmelCase=4096			,       __lowerCAmelCase=12			,       __lowerCAmelCase=1			,       __lowerCAmelCase=64			,       __lowerCAmelCase=16384			,       __lowerCAmelCase=1			,       __lowerCAmelCase="gelu_new"			,       __lowerCAmelCase=0			,       __lowerCAmelCase=0			,       __lowerCAmelCase=512			,       __lowerCAmelCase=2			,       __lowerCAmelCase=0.02			,       __lowerCAmelCase=1E-1_2			,       __lowerCAmelCase=0.1			,       __lowerCAmelCase="absolute"			,       __lowerCAmelCase=0			,       __lowerCAmelCase=2			,       __lowerCAmelCase=3			,       **__lowerCAmelCase			,       ):




					"""simple docstring"""


					super().__init__(pad_token_id=__lowerCAmelCase			,       bos_token_id=__lowerCAmelCase			,       eos_token_id=__lowerCAmelCase			,       **__lowerCAmelCase)

					lowerCAmelCase     =    vocab_size
					lowerCAmelCase     =    embedding_size
					lowerCAmelCase     =    hidden_size
					lowerCAmelCase     =    num_hidden_layers
					lowerCAmelCase     =    num_hidden_groups
					lowerCAmelCase     =    num_attention_heads
					lowerCAmelCase     =    inner_group_num
					lowerCAmelCase     =    hidden_act
					lowerCAmelCase     =    intermediate_size
					lowerCAmelCase     =    hidden_dropout_prob
					lowerCAmelCase     =    attention_probs_dropout_prob
					lowerCAmelCase     =    max_position_embeddings
					lowerCAmelCase     =    type_vocab_size
					lowerCAmelCase     =    initializer_range
					lowerCAmelCase     =    layer_norm_eps
					lowerCAmelCase     =    classifier_dropout_prob
					lowerCAmelCase     =    position_embedding_type







class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			@property
			def 				a_			(  self):




					"""simple docstring"""


					if self.task == "multiple-choice":
							lowerCAmelCase     =    {0: """batch""", 1: """choice""", 2: """sequence"""}
					else:
							lowerCAmelCase     =    {0: """batch""", 1: """sequence"""}
					return OrderedDict(
					    [
					        ("""input_ids""", dynamic_axis),
					        ("""attention_mask""", dynamic_axis),
					        ("""token_type_ids""", dynamic_axis),
					    ])




