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


UpperCamelCase			:       Union[str, 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 __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





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

      '''simple docstring'''




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





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None   ):

      '''simple docstring'''




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

      __UpperCamelCase       						= self._remove_unused_columns(__UpperCAmelCase			,  description='Calibration'   )

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





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.train_dataset if calib_dataset is None else calib_dataset
      __UpperCamelCase       						= self.get_calib_dataloader(__UpperCAmelCase   )

      __UpperCamelCase       						= self.model
      quant_trainer.configure_model(__UpperCAmelCase			,  self.quant_trainer_args			,  calib=__UpperCAmelCase   )
      model.eval()
      quant_trainer.enable_calibration(__UpperCAmelCase   )

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

      quant_trainer.finish_calibration(__UpperCAmelCase			,  self.quant_trainer_args   )
      __UpperCamelCase       						= model





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase = "eval"   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.eval_dataset if eval_dataset is None else eval_dataset
      __UpperCamelCase       						= self.get_eval_dataloader(__UpperCAmelCase   )
      __UpperCamelCase       						= self.eval_examples if eval_examples is None else eval_examples

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

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

       # Prefix all keys with metric_key_prefix + '_'
       for key in list(metrics.keys()   ):
        if not key.startswith(F'{metric_key_prefix}_'   ):
         __UpperCamelCase       						= metrics.pop(__UpperCAmelCase   )

       self.log(__UpperCAmelCase   )
      else:
       __UpperCamelCase       						= {}

      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       						= self.callback_handler.on_evaluate(self.args			,  self.state			,  self.control			,  __UpperCAmelCase   )
      return metrics





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase=None			,  __UpperCAmelCase = "test"   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.get_test_dataloader(__UpperCAmelCase   )

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

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

      __UpperCamelCase       						= self.post_process_function(__UpperCAmelCase			,  __UpperCAmelCase			,  output.predictions			,  'predict'   )
      __UpperCamelCase       						= self.compute_metrics(__UpperCAmelCase   )

      # Prefix all keys with metric_key_prefix + '_'
      for key in list(metrics.keys()   ):
       if not key.startswith(F'{metric_key_prefix}_'   ):
        __UpperCamelCase       						= metrics.pop(__UpperCAmelCase   )

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







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase="./"   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.eval_dataset
      __UpperCamelCase       						= self.get_eval_dataloader(__UpperCAmelCase   )

      __UpperCamelCase       						= next(iter(__UpperCAmelCase   )   )

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

      # convert to tuple
      __UpperCamelCase       						= tuple(v.to(__UpperCAmelCase   ) for k, v in batch.items()   )

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

      __UpperCamelCase       						= True

      __UpperCamelCase       						= self.model.to(__UpperCAmelCase   )

      model.eval()
      model.float()

      __UpperCamelCase       						= model.module if hasattr(__UpperCAmelCase			,  'module'   ) else model
      quant_trainer.configure_model(__UpperCAmelCase			,  self.quant_trainer_args   )

      __UpperCamelCase       						= os.path.join(__UpperCAmelCase			,  'model.onnx'   )
      logger.info(F'exporting model to {output_model_file}'   )

      __UpperCamelCase       						= {0: 'batch_size', 1: 'seq_len'}

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





"""simple docstring"""

from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCamelCase			:       List[str]									=				{
    "configuration_autoformer": [
        "AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "AutoformerConfig",
    ],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCamelCase			:       Dict									=				[
          "AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
          "AutoformerForPrediction",
          "AutoformerModel",
          "AutoformerPreTrainedModel",
      ]


if TYPE_CHECKING:
      from .configuration_autoformer import (
          AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
          AutoformerConfig,
      )

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_autoformer import (
                AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                AutoformerForPrediction,
                AutoformerModel,
                AutoformerPreTrainedModel,
            )

else:
      import sys

      UpperCamelCase			:       int									=				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



