





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__:      int									=		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     SCREAMING_SNAKE_CASE__ (						UpperCamelCase__      ):



 def __init__(							self,*__lowerCamelCase,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,**__lowerCamelCase ):
        super().__init__(*__lowerCamelCase,**__lowerCamelCase )
        A__											=						eval_examples
        A__											=						post_process_function
        A__											=						quant_trainer_args
        A__											=						128  # default number of calibration samples



 def 				UpperCamelCase			(							self,__lowerCamelCase=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(__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 				UpperCamelCase			(							self,__lowerCamelCase=None ):
        A__											=						self.train_dataset if calib_dataset is None else calib_dataset
        A__											=						self.get_calib_dataloader(__lowerCamelCase )

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



 def 				UpperCamelCase			(							self,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase=None,__lowerCamelCase = "eval" ):
        A__											=						self.eval_dataset if eval_dataset is None else eval_dataset
        A__											=						self.get_eval_dataloader(__lowerCamelCase )
        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(
                   __lowerCamelCase,description='''Evaluation''',prediction_loss_only=True if compute_metrics is None else None,ignore_keys=__lowerCamelCase,)
        finally:
               A__											=						compute_metrics

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

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



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase=None,__lowerCamelCase = "test" ):
        A__											=						self.get_test_dataloader(__lowerCamelCase )

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

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

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

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




 def 				UpperCamelCase			(							self,__lowerCamelCase="./" ):
        A__											=						self.eval_dataset
        A__											=						self.get_eval_dataloader(__lowerCamelCase )

        A__											=						next(iter(__lowerCamelCase ) )

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

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

        model.eval()
        model.float()

        A__											=						model.module if hasattr(__lowerCamelCase,'''module''' ) else model
        quant_trainer.configure_model(__lowerCamelCase,self.quant_trainer_args )

        A__											=						os.path.join(__lowerCamelCase,'''model.onnx''' )
        logger.info(f"exporting model to {output_model_file}" )

        A__											=						{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''' )












class     SCREAMING_SNAKE_CASE__ :



 def __init__(							self,__lowerCamelCase ):
        A__											=						set_counts
        A__											=						max(__lowerCamelCase )
        A__											=						len(__lowerCamelCase )
        A__											=						[1] * num_sets
        A__											=						list(range(__lowerCamelCase ) )



 def 				UpperCamelCase			(							self,__lowerCamelCase,__lowerCamelCase ):
        A__											=						self.get_parent(__lowerCamelCase )
        A__											=						self.get_parent(__lowerCamelCase )

        if src_parent == dst_parent:
               return False

        if self.ranks[dst_parent] >= self.ranks[src_parent]:
               self.set_counts[dst_parent] += self.set_counts[src_parent]
               A__											=						0
               A__											=						dst_parent
               if self.ranks[dst_parent] == self.ranks[src_parent]:
                      self.ranks[dst_parent] += 1
               A__											=						self.set_counts[dst_parent]
        else:
               self.set_counts[src_parent] += self.set_counts[dst_parent]
               A__											=						0
               A__											=						src_parent
               A__											=						self.set_counts[src_parent]

        A__											=						max(self.max_set,__lowerCamelCase )
        return True




 def 				UpperCamelCase			(							self,__lowerCamelCase ):
        if self.parents[disj_set] == disj_set:
               return disj_set
        A__											=						self.get_parent(self.parents[disj_set] )
        return self.parents[disj_set]





