






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






 """simple docstring"""





 def __init__(							self			,     *__lowerCAmelCase			,     __lowerCAmelCase=None			,     __lowerCAmelCase=None			,     __lowerCAmelCase=None			,     **__lowerCAmelCase ):
       super().__init__(*__lowerCAmelCase			,     **__lowerCAmelCase )
       UpperCamelCase_				:   List[str]					    =							eval_examples
       UpperCamelCase_				:   Optional[Any]					    =							post_process_function
       UpperCamelCase_				:   Tuple					    =							quant_trainer_args
       UpperCamelCase_				:   Optional[Any]					    =							1_28  # 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.""" )
       UpperCamelCase_				:   Dict					    =							calib_dataset if calib_dataset is not None else self.calib_dataset

       UpperCamelCase_				:   Any					    =							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 ):
       UpperCamelCase_				:   Union[str, Any]					    =							self.train_dataset if calib_dataset is None else calib_dataset
       UpperCamelCase_				:   str					    =							self.get_calib_dataloader(__lowerCAmelCase )

       UpperCamelCase_				:   int					    =							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
             UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_				:   int					    =							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 )
       UpperCamelCase_				:   List[Any]					    =							model
 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase=None			,     __lowerCAmelCase=None			,     __lowerCAmelCase=None			,     __lowerCAmelCase = "eval" ):
       UpperCamelCase_				:   List[str]					    =							self.eval_dataset if eval_dataset is None else eval_dataset
       UpperCamelCase_				:   Tuple					    =							self.get_eval_dataloader(__lowerCAmelCase )
       UpperCamelCase_				:   Tuple					    =							self.eval_examples if eval_examples is None else eval_examples

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

       if self.post_process_function is not None and self.compute_metrics is not None:
             UpperCamelCase_				:   Optional[Any]					    =							self.post_process_function(__lowerCAmelCase			,     __lowerCAmelCase			,     output.predictions )
             UpperCamelCase_				:   Tuple					    =							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}_" ):
                         UpperCamelCase_				:   Union[str, Any]					    =							metrics.pop(__lowerCAmelCase )

             self.log(__lowerCAmelCase )
       else:
             UpperCamelCase_				:   Optional[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() )

       UpperCamelCase_				:   int					    =							self.callback_handler.on_evaluate(self.args			,     self.state			,     self.control			,     __lowerCAmelCase )
       return metrics
 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase=None			,     __lowerCAmelCase = "test" ):
       UpperCamelCase_				:   int					    =							self.get_test_dataloader(__lowerCAmelCase )

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

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

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

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





 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase="./" ):
       UpperCamelCase_				:   Any					    =							self.eval_dataset
       UpperCamelCase_				:   Union[str, Any]					    =							self.get_eval_dataloader(__lowerCAmelCase )

       UpperCamelCase_				:   Optional[int]					    =							next(iter(__lowerCAmelCase ) )

       # saving device - to make it consistent
       UpperCamelCase_				:   List[str]					    =							torch.device("""cuda""" if torch.cuda.is_available() else """cpu""" )

       # convert to tuple
       UpperCamelCase_				:   Dict					    =							tuple(v.to(__lowerCAmelCase ) for k, v in batch.items() )

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

       UpperCamelCase_				:   Optional[int]					    =							True

       UpperCamelCase_				:   Tuple					    =							self.model.to(__lowerCAmelCase )

       model.eval()
       model.float()

       UpperCamelCase_				:   Optional[int]					    =							model.module if hasattr(__lowerCAmelCase			,     """module""" ) else model
       quant_trainer.configure_model(__lowerCAmelCase			,     self.quant_trainer_args )

       UpperCamelCase_				:   List[str]					    =							os.path.join(__lowerCAmelCase			,     """model.onnx""" )
       logger.info(F"exporting model to {output_model_file}" )

       UpperCamelCase_				:   Optional[Any]					    =							{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""" )














'''simple docstring'''

# Lint as: python3
import itertools
import os
import re


UpperCamelCase							=re.compile(R"([A-Z]+)([A-Z][a-z])")
UpperCamelCase							=re.compile(R"([a-z\d])([A-Z])")

UpperCamelCase							=re.compile(R"(?<!_)_(?!_)")
UpperCamelCase							=re.compile(R"(_{2,})")

UpperCamelCase							=R"^\w+(\.\w+)*$"

UpperCamelCase							=R"<>:/\|?*"



def 							snake_case					(     a_							: str	)       ->				int:





      """simple docstring"""

      UpperCamelCase_				:   str					    =							_uppercase_uppercase_re.sub(r"""\1_\2"""    ,   a_	)
      UpperCamelCase_				:   List[Any]					    =							_lowercase_uppercase_re.sub(r"""\1_\2"""    ,   a_	)
      return name.lower()



def 							snake_case					(     a_							: Any	)       ->				Optional[Any]:





      """simple docstring"""

      UpperCamelCase_				:   Optional[Any]					    =							_single_underscore_re.split(a_	)
      UpperCamelCase_				:   Tuple					    =							[_multiple_underscores_re.split(a_	) for n in name]
      return "".join(n.capitalize() for n in itertools.chain.from_iterable(a_	) if n != """"""	)



def 							snake_case					(     a_							: Optional[int]	)       ->				str:





      """simple docstring"""

      if os.path.basename(a_	) != name:
            raise ValueError(f"Should be a dataset name, not a path: {name}"	)
      return camelcase_to_snakecase(a_	)



def 							snake_case					(     a_							: List[Any]    ,   a_							: Optional[int]	)       ->				Optional[Any]:





      """simple docstring"""

      if os.path.basename(a_	) != name:
            raise ValueError(f"Should be a dataset name, not a path: {name}"	)
      if not re.match(_split_re    ,   a_	):
            raise ValueError(f"Split name should match '{_split_re}'' but got '{split}'."	)
      return f"{filename_prefix_for_name(a_	)}-{split}"



def 							snake_case					(     a_							: Optional[Any]    ,   a_							: int    ,   a_							: List[Any]    ,   a_							: Union[str, Any]=None	)       ->				Dict:





      """simple docstring"""

      UpperCamelCase_				:   Optional[int]					    =							filename_prefix_for_split(a_    ,   a_	)
      if filetype_suffix:
            prefix += f".{filetype_suffix}"
      UpperCamelCase_				:   str					    =							os.path.join(a_    ,   a_	)
      return f"{filepath}*"



def 							snake_case					(     a_							: Tuple    ,   a_							: Any    ,   a_							: Optional[int]    ,   a_							: Union[str, Any]=None    ,   a_							: Optional[Any]=None	)       ->				List[Any]:





      """simple docstring"""

      UpperCamelCase_				:   Dict					    =							filename_prefix_for_split(a_    ,   a_	)
      UpperCamelCase_				:   Dict					    =							os.path.join(a_    ,   a_	)

      if shard_lengths:
            UpperCamelCase_				:   Optional[int]					    =							len(a_	)
            UpperCamelCase_				:   Tuple					    =							[f"{prefix}-{shard_id:05d}-of-{num_shards:05d}" for shard_id in range(a_	)]
            if filetype_suffix:
                  UpperCamelCase_				:   Tuple					    =							[filename + f".{filetype_suffix}" for filename in filenames]
            return filenames
      else:
            UpperCamelCase_				:   Dict					    =							prefix
            if filetype_suffix:
                  filename += f".{filetype_suffix}"
            return [filename]






