


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   							=						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			(_a						):







       def __init__(			self: List[str],*A_: Optional[Any],A_: Tuple=None,A_: str=None,A_: str=None,**A_: List[Any]					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: int,A_: Union[str, Any]=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(A_,description='Calibration'					)

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







       def 		snake_case_	(			self: Any,A_: Tuple=None					):






            '''simple docstring'''




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

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

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

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







       def 		snake_case_	(			self: List[str],A_: str=None,A_: int=None,A_: str=None,A_: str = "eval"					):






            '''simple docstring'''




            __UpperCamelCase =			self.eval_dataset if eval_dataset is None else eval_dataset
            __UpperCamelCase =			self.get_eval_dataloader(A_					)
            __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(
                     A_,description='Evaluation',prediction_loss_only=True if compute_metrics is None else None,ignore_keys=A_,)
            finally:
                 __UpperCamelCase =			compute_metrics

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

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

                 self.log(A_					)
            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,A_					)
            return metrics







       def 		snake_case_	(			self: List[Any],A_: Optional[int],A_: Tuple,A_: Any=None,A_: str = "test"					):






            '''simple docstring'''




            __UpperCamelCase =			self.get_test_dataloader(A_					)

            # 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(
                     A_,description='Prediction',prediction_loss_only=True if compute_metrics is None else None,ignore_keys=A_,)
            finally:
                 __UpperCamelCase =			compute_metrics

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

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

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

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







       def 		snake_case_	(			self: Dict,A_: Optional[Any]="./"					):






            '''simple docstring'''




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

            __UpperCamelCase =			next(iter(A_					)					)

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

            # convert to tuple
            __UpperCamelCase =			tuple(v.to(A_					) 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(A_					)

            model.eval()
            model.float()

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

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

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

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






def 		_A		(		_lowercase )					-> int:


     """simple docstring"""


     assert column_title.isupper()
     __UpperCamelCase =			0
     __UpperCamelCase =			len(_lowercase ) - 1
     __UpperCamelCase =			0

     while index >= 0:
          __UpperCamelCase =			(ord(column_title[index] ) - 64) * pow(26					,  _lowercase )
          answer += value
          power += 1
          index -= 1

     return answer


if __name__ == "__main__":
    from doctest import testmod

    testmod()


