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


SCREAMING_SNAKE_CASE :	Optional[Any]							     =		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_ (       a_ ):

  def __init__(		self :      Optional[Any]     ,							*__SCREAMING_SNAKE_CASE :      int     ,							__SCREAMING_SNAKE_CASE :      str=None     ,							__SCREAMING_SNAKE_CASE :      Tuple=None     ,							__SCREAMING_SNAKE_CASE :      Optional[int]=None     ,							**__SCREAMING_SNAKE_CASE :      Tuple       ):
   super().__init__(*__SCREAMING_SNAKE_CASE     ,							**__SCREAMING_SNAKE_CASE       )
   __a       =  eval_examples
   __a       =  post_process_function
   __a       =  quant_trainer_args
   __a       =  1_28  # default number of calibration samples

  def 		_UpperCAmelCase   (		self :      str     ,							__SCREAMING_SNAKE_CASE :      str=None       ):
   if calib_dataset is None and self.calib_dataset is None:
    raise ValueError("Trainer: calibration requires an calib_dataset."       )
   __a       =  calib_dataset if calib_dataset is not None else self.calib_dataset

   __a       =  self._remove_unused_columns(__SCREAMING_SNAKE_CASE     ,							description="Calibration"       )

   return DataLoader(
       __SCREAMING_SNAKE_CASE     ,							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=__SCREAMING_SNAKE_CASE     ,							)

  def 		_UpperCAmelCase   (		self :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      Tuple=None       ):
   __a       =  self.train_dataset if calib_dataset is None else calib_dataset
   __a       =  self.get_calib_dataloader(__SCREAMING_SNAKE_CASE       )

   __a       =  self.model
   quant_trainer.configure_model(__SCREAMING_SNAKE_CASE     ,							self.quant_trainer_args     ,							calib=__SCREAMING_SNAKE_CASE       )
   model.eval()
   quant_trainer.enable_calibration(__SCREAMING_SNAKE_CASE       )

   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(__SCREAMING_SNAKE_CASE       ):
    # Prediction step
    __a       ,     __a       ,     __a       =  self.prediction_step(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							prediction_loss_only=__SCREAMING_SNAKE_CASE       )
    if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
     break

   quant_trainer.finish_calibration(__SCREAMING_SNAKE_CASE     ,							self.quant_trainer_args       )
   __a       =  model

  def 		_UpperCAmelCase   (		self :      List[Any]     ,							__SCREAMING_SNAKE_CASE :      List[str]=None     ,							__SCREAMING_SNAKE_CASE :      str=None     ,							__SCREAMING_SNAKE_CASE :      Tuple=None     ,							__SCREAMING_SNAKE_CASE :      str = "eval"       ):
   __a       =  self.eval_dataset if eval_dataset is None else eval_dataset
   __a       =  self.get_eval_dataloader(__SCREAMING_SNAKE_CASE       )
   __a       =  self.eval_examples if eval_examples is None else eval_examples

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

   if self.post_process_function is not None and self.compute_metrics is not None:
    __a       =  self.post_process_function(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							output.predictions       )
    __a       =  self.compute_metrics(__SCREAMING_SNAKE_CASE       )

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

    self.log(__SCREAMING_SNAKE_CASE       )
   else:
    __a       =  {}

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

   __a       =  self.callback_handler.on_evaluate(self.args     ,							self.state     ,							self.control     ,							__SCREAMING_SNAKE_CASE       )
   return metrics

  def 		_UpperCAmelCase   (		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      List[str]     ,							__SCREAMING_SNAKE_CASE :      Dict     ,							__SCREAMING_SNAKE_CASE :      List[Any]=None     ,							__SCREAMING_SNAKE_CASE :      str = "test"       ):
   __a       =  self.get_test_dataloader(__SCREAMING_SNAKE_CASE       )

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

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

   __a       =  self.post_process_function(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							output.predictions     ,							"predict"       )
   __a       =  self.compute_metrics(__SCREAMING_SNAKE_CASE       )

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

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

  def 		_UpperCAmelCase   (		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      List[str]="./"       ):
   __a       =  self.eval_dataset
   __a       =  self.get_eval_dataloader(__SCREAMING_SNAKE_CASE       )

   __a       =  next(iter(__SCREAMING_SNAKE_CASE       )       )

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

   # convert to tuple
   __a       =  tuple(v.to(__SCREAMING_SNAKE_CASE       ) for k, v in batch.items()       )

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

   __a       =  True

   __a       =  self.model.to(__SCREAMING_SNAKE_CASE       )

   model.eval()
   model.float()

   __a       =  model.module if hasattr(__SCREAMING_SNAKE_CASE     ,							"module"       ) else model
   quant_trainer.configure_model(__SCREAMING_SNAKE_CASE     ,							self.quant_trainer_args       )

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

   __a       =  {0: "batch_size", 1: "seq_len"}

   torch.onnx.export(
       __SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							export_params=__SCREAMING_SNAKE_CASE     ,							opset_version=13     ,							do_constant_folding=__SCREAMING_SNAKE_CASE     ,							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=__SCREAMING_SNAKE_CASE     ,							)
   logger.info("onnx export finished"       )



from ... import PretrainedConfig


SCREAMING_SNAKE_CASE :	Any							     =		{
    """sijunhe/nezha-cn-base""": """https://huggingface.co/sijunhe/nezha-cn-base/resolve/main/config.json""",
}
class    A_ (       a_ ):
  _SCREAMING_SNAKE_CASE    				=			NEZHA_PRETRAINED_CONFIG_ARCHIVE_MAP
  _SCREAMING_SNAKE_CASE    				=			"""nezha"""

  def __init__(		self :      Any     ,							__SCREAMING_SNAKE_CASE :      List[str]=2_11_28     ,							__SCREAMING_SNAKE_CASE :      Dict=7_68     ,							__SCREAMING_SNAKE_CASE :      List[Any]=12     ,							__SCREAMING_SNAKE_CASE :      str=12     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=30_72     ,							__SCREAMING_SNAKE_CASE :      Any="gelu"     ,							__SCREAMING_SNAKE_CASE :      int=0.1     ,							__SCREAMING_SNAKE_CASE :      List[str]=0.1     ,							__SCREAMING_SNAKE_CASE :      int=5_12     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]=64     ,							__SCREAMING_SNAKE_CASE :      List[Any]=2     ,							__SCREAMING_SNAKE_CASE :      int=0.02     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=1E-12     ,							__SCREAMING_SNAKE_CASE :      int=0.1     ,							__SCREAMING_SNAKE_CASE :      Any=0     ,							__SCREAMING_SNAKE_CASE :      Any=2     ,							__SCREAMING_SNAKE_CASE :      Optional[Any]=3     ,							__SCREAMING_SNAKE_CASE :      List[Any]=True     ,							**__SCREAMING_SNAKE_CASE :      List[str]     ,							):
   super().__init__(pad_token_id=__SCREAMING_SNAKE_CASE     ,							bos_token_id=__SCREAMING_SNAKE_CASE     ,							eos_token_id=__SCREAMING_SNAKE_CASE     ,							**__SCREAMING_SNAKE_CASE       )

   __a       =  vocab_size
   __a       =  hidden_size
   __a       =  num_hidden_layers
   __a       =  num_attention_heads
   __a       =  hidden_act
   __a       =  intermediate_size
   __a       =  hidden_dropout_prob
   __a       =  attention_probs_dropout_prob
   __a       =  max_position_embeddings
   __a       =  max_relative_position
   __a       =  type_vocab_size
   __a       =  initializer_range
   __a       =  layer_norm_eps
   __a       =  classifier_dropout
   __a       =  use_cache


