
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


_snake_case      :							Tuple     	=    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   UpperCamelCase_      (      __a		):







						'''simple docstring'''



						def __init__(			self							:Any		,    *lowerCAmelCase__							:str		,    lowerCAmelCase__							:Optional[int]=None		,    lowerCAmelCase__							:Any=None		,    lowerCAmelCase__							:str=None		,    **lowerCAmelCase__							:Any     )  ->List[Any]:
										super().__init__(*lowerCAmelCase__		,    **lowerCAmelCase__     )
										lowercase         =      eval_examples
										lowercase         =      post_process_function
										lowercase         =      quant_trainer_args
										lowercase         =      128  # default number of calibration samples





						def        SCREAMING_SNAKE_CASE(			self							:Any		,    lowerCAmelCase__							:Optional[Any]=None     )  ->Tuple:
										if calib_dataset is None and self.calib_dataset is None:
														raise ValueError("Trainer: calibration requires an calib_dataset."     )
										lowercase         =      calib_dataset if calib_dataset is not None else self.calib_dataset

										lowercase         =      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        SCREAMING_SNAKE_CASE(			self							:Optional[int]		,    lowerCAmelCase__							:Union[str, Any]=None     )  ->Any:
										lowercase         =      self.train_dataset if calib_dataset is None else calib_dataset
										lowercase         =      self.get_calib_dataloader(lowerCAmelCase__     )

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





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]		,    lowerCAmelCase__							:str=None		,    lowerCAmelCase__							:Dict=None		,    lowerCAmelCase__							:Dict=None		,    lowerCAmelCase__							:str = "eval"     )  ->str:
										lowercase         =      self.eval_dataset if eval_dataset is None else eval_dataset
										lowercase         =      self.get_eval_dataloader(lowerCAmelCase__     )
										lowercase         =      self.eval_examples if eval_examples is None else eval_examples

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

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

														self.log(lowerCAmelCase__     )
										else:
														lowercase         =      {}

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

										lowercase         =      self.callback_handler.on_evaluate(self.args		,    self.state		,    self.control		,    lowerCAmelCase__     )
										return metrics





						def        SCREAMING_SNAKE_CASE(			self							:int		,    lowerCAmelCase__							:str		,    lowerCAmelCase__							:Union[str, Any]		,    lowerCAmelCase__							:Dict=None		,    lowerCAmelCase__							:str = "test"     )  ->Dict:
										lowercase         =      self.get_test_dataloader(lowerCAmelCase__     )

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

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

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

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





						def        SCREAMING_SNAKE_CASE(			self							:Any		,    lowerCAmelCase__							:Union[str, Any]="./"     )  ->Any:
										lowercase         =      self.eval_dataset
										lowercase         =      self.get_eval_dataloader(lowerCAmelCase__     )

										lowercase         =      next(iter(lowerCAmelCase__     )     )

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

										# convert to tuple
										lowercase         =      tuple(v.to(lowerCAmelCase__     ) for k, v in batch.items()     )

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

										lowercase         =      True

										lowercase         =      self.model.to(lowerCAmelCase__     )

										model.eval()
										model.float()

										lowercase         =      model.module if hasattr(lowerCAmelCase__		,    "module"     ) else model
										quant_trainer.configure_model(lowerCAmelCase__		,    self.quant_trainer_args     )

										lowercase         =      os.path.join(lowerCAmelCase__		,    "model.onnx"     )
										logger.info(F'''exporting model to {output_model_file}'''     )

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


from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

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


if TYPE_CHECKING:
						from ... import FeatureExtractionMixin, PreTrainedTokenizerBase, TensorType


_snake_case      :							Dict     	=    logging.get_logger(__name__)

_snake_case      :							List[str]     	=    {
    "microsoft/deberta-v2-xlarge": "https://huggingface.co/microsoft/deberta-v2-xlarge/resolve/main/config.json",
    "microsoft/deberta-v2-xxlarge": "https://huggingface.co/microsoft/deberta-v2-xxlarge/resolve/main/config.json",
    "microsoft/deberta-v2-xlarge-mnli": (
        "https://huggingface.co/microsoft/deberta-v2-xlarge-mnli/resolve/main/config.json"
    ),
    "microsoft/deberta-v2-xxlarge-mnli": (
        "https://huggingface.co/microsoft/deberta-v2-xxlarge-mnli/resolve/main/config.json"
    ),
}





class   UpperCamelCase_      (      __a		):







						'''simple docstring'''



						UpperCamelCase      :  Optional[Any]				   =					'''deberta-v2'''





						def __init__(			self							:int		,    lowerCAmelCase__							:Union[str, Any]=128100		,    lowerCAmelCase__							:Optional[Any]=1536		,    lowerCAmelCase__							:Tuple=24		,    lowerCAmelCase__							:List[Any]=24		,    lowerCAmelCase__							:Union[str, Any]=6144		,    lowerCAmelCase__							:List[str]="gelu"		,    lowerCAmelCase__							:Union[str, Any]=0.1		,    lowerCAmelCase__							:Optional[Any]=0.1		,    lowerCAmelCase__							:List[str]=512		,    lowerCAmelCase__							:Optional[int]=0		,    lowerCAmelCase__							:Optional[int]=0.02		,    lowerCAmelCase__							:int=1E-7		,    lowerCAmelCase__							:Optional[int]=False		,    lowerCAmelCase__							:List[Any]=-1		,    lowerCAmelCase__							:str=0		,    lowerCAmelCase__							:str=True		,    lowerCAmelCase__							:Optional[Any]=None		,    lowerCAmelCase__							:List[Any]=0		,    lowerCAmelCase__							:Union[str, Any]="gelu"		,    **lowerCAmelCase__							:str		,    )  ->Optional[Any]:
										super().__init__(**lowerCAmelCase__     )

										lowercase         =      hidden_size
										lowercase         =      num_hidden_layers
										lowercase         =      num_attention_heads
										lowercase         =      intermediate_size
										lowercase         =      hidden_act
										lowercase         =      hidden_dropout_prob
										lowercase         =      attention_probs_dropout_prob
										lowercase         =      max_position_embeddings
										lowercase         =      type_vocab_size
										lowercase         =      initializer_range
										lowercase         =      relative_attention
										lowercase         =      max_relative_positions
										lowercase         =      pad_token_id
										lowercase         =      position_biased_input

										# Backwards compatibility
										if type(lowerCAmelCase__     ) == str:
														lowercase         =      [x.strip() for x in pos_att_type.lower().split("|"     )]

										lowercase         =      pos_att_type
										lowercase         =      vocab_size
										lowercase         =      layer_norm_eps

										lowercase         =      kwargs.get("pooler_hidden_size"		,    lowerCAmelCase__     )
										lowercase         =      pooler_dropout
										lowercase         =      pooler_hidden_act





class   UpperCamelCase_      (      __a		):







						'''simple docstring'''



						@property
						def        SCREAMING_SNAKE_CASE(			self							:str     )  ->Mapping[str, Mapping[int, str]]:
										if self.task == "multiple-choice":
														lowercase         =      {0: "batch", 1: "choice", 2: "sequence"}
										else:
														lowercase         =      {0: "batch", 1: "sequence"}
										if self._config.type_vocab_size > 0:
														return OrderedDict(
														    [("input_ids", dynamic_axis), ("attention_mask", dynamic_axis), ("token_type_ids", dynamic_axis)]     )
										else:
														return OrderedDict([("input_ids", dynamic_axis), ("attention_mask", dynamic_axis)]     )





						@property
						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->int:
										return 12





						def        SCREAMING_SNAKE_CASE(			self							:Union[str, Any]		,    lowerCAmelCase__							:Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]		,    lowerCAmelCase__							:int = -1		,    lowerCAmelCase__							:int = -1		,    lowerCAmelCase__							:int = -1		,    lowerCAmelCase__							:bool = False		,    lowerCAmelCase__							:Optional["TensorType"] = None		,    lowerCAmelCase__							:int = 3		,    lowerCAmelCase__							:int = 40		,    lowerCAmelCase__							:int = 40		,    lowerCAmelCase__							:"PreTrainedTokenizerBase" = None		,    )  ->Mapping[str, Any]:
										lowercase         =      super().generate_dummy_inputs(preprocessor=lowerCAmelCase__		,    framework=lowerCAmelCase__     )
										if self._config.type_vocab_size == 0 and "token_type_ids" in dummy_inputs:
														del dummy_inputs["token_type_ids"]
										return dummy_inputs
