





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


__A											=						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			):
       """simple docstring"""




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



         '''simple docstring'''

         super().__init__(*__UpperCAmelCase       ,					**__UpperCAmelCase )
         lowerCAmelCase__     :Tuple											=				eval_examples
         lowerCAmelCase__     :Optional[Any]											=				post_process_function
         lowerCAmelCase__     :Tuple											=				quant_trainer_args
         lowerCAmelCase__     :Tuple											=				1_2_8  # default number of calibration samples






       def 	snake_case		(       self       ,					__UpperCAmelCase=None ):



         '''simple docstring'''

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

         lowerCAmelCase__     :Optional[int]											=				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 	snake_case		(       self       ,					__UpperCAmelCase=None ):



         '''simple docstring'''

         lowerCAmelCase__     :List[Any]											=				self.train_dataset if calib_dataset is None else calib_dataset
         lowerCAmelCase__     :Tuple											=				self.get_calib_dataloader(__UpperCAmelCase )

         lowerCAmelCase__     :List[str]											=				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
           lowerCAmelCase__				,							lowerCAmelCase__				,							lowerCAmelCase__     :Tuple											=				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 )
         lowerCAmelCase__     :List[str]											=				model






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



         '''simple docstring'''

         lowerCAmelCase__     :Union[str, Any]											=				self.eval_dataset if eval_dataset is None else eval_dataset
         lowerCAmelCase__     :List[Any]											=				self.get_eval_dataloader(__UpperCAmelCase )
         lowerCAmelCase__     :List[str]											=				self.eval_examples if eval_examples is None else eval_examples

         # Temporarily disable metric computation, we will do it in the loop here.
         lowerCAmelCase__     :Union[str, Any]											=				self.compute_metrics
         lowerCAmelCase__     :List[Any]											=				None
         lowerCAmelCase__     :Optional[int]											=				self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
         try:
           lowerCAmelCase__     :Any											=				eval_loop(
               __UpperCAmelCase       ,					description='Evaluation'       ,					prediction_loss_only=True if compute_metrics is None else None       ,					ignore_keys=__UpperCAmelCase       ,					)
         finally:
           lowerCAmelCase__     :int											=				compute_metrics

         if self.post_process_function is not None and self.compute_metrics is not None:
           lowerCAmelCase__     :List[Any]											=				self.post_process_function(__UpperCAmelCase       ,					__UpperCAmelCase       ,					output.predictions )
           lowerCAmelCase__     :Any											=				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}_" ):
               lowerCAmelCase__     :List[Any]											=				metrics.pop(__UpperCAmelCase )

           self.log(__UpperCAmelCase )
         else:
           lowerCAmelCase__     :int											=				{}

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






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



         '''simple docstring'''

         lowerCAmelCase__     :List[Any]											=				self.get_test_dataloader(__UpperCAmelCase )

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

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

         lowerCAmelCase__     :List[Any]											=				self.post_process_function(__UpperCAmelCase       ,					__UpperCAmelCase       ,					output.predictions       ,					'predict' )
         lowerCAmelCase__     :Dict											=				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}_" ):
             lowerCAmelCase__     :Dict											=				metrics.pop(__UpperCAmelCase )

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






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



         '''simple docstring'''

         lowerCAmelCase__     :str											=				self.eval_dataset
         lowerCAmelCase__     :int											=				self.get_eval_dataloader(__UpperCAmelCase )

         lowerCAmelCase__     :List[str]											=				next(iter(__UpperCAmelCase ) )

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

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

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

         lowerCAmelCase__     :int											=				True

         lowerCAmelCase__     :Optional[Any]											=				self.model.to(__UpperCAmelCase )

         model.eval()
         model.float()

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

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

         lowerCAmelCase__     :Any											=				{0: 'batch_size', 1: 'seq_len'}

         torch.onnx.export(
             __UpperCAmelCase       ,					__UpperCAmelCase       ,					__UpperCAmelCase       ,					export_params=__UpperCAmelCase       ,					opset_version=1_3       ,					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 math import loga




def 						__lowerCamelCase			(	UpperCAmelCase_							)   ->int:
	if a < 0:
		raise ValueError('Input value must be a positive integer'							)
	elif isinstance(UpperCAmelCase_					,				UpperCAmelCase_							):
		raise TypeError('Input value must be a \'int\' type'							)
	return 0 if (a == 0) else int(loga(a & -a							)							)


if __name__ == "__main__":
			import doctest

			doctest.testmod()



