

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


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



    def __init__(       self		,					*UpperCamelCase__		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					**UpperCamelCase__							)    ->   Optional[int]:





           '''simple docstring'''
           super().__init__(*UpperCamelCase__		,					**UpperCamelCase__							)
           A_						       =   eval_examples
           A_						       =   post_process_function
           A_						       =   quant_trainer_args
           A_						       =   128  # default number of calibration samples



    def        snake_case_      (       self		,					UpperCamelCase__=None							)    ->   Union[str, Any]:





           '''simple docstring'''

           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(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							)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   self.train_dataset if calib_dataset is None else calib_dataset
           A_						       =   self.get_calib_dataloader(UpperCamelCase__							)

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



    def        snake_case_      (       self		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					UpperCamelCase__ = "eval"							)    ->   Optional[Any]:





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

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

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



    def        snake_case_      (       self		,					UpperCamelCase__		,					UpperCamelCase__		,					UpperCamelCase__=None		,					UpperCamelCase__ = "test"							)    ->   Union[str, Any]:





           '''simple docstring'''
           A_						       =   self.get_test_dataloader(UpperCamelCase__							)

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

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

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

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







    def        snake_case_      (       self		,					UpperCamelCase__="./"							)    ->   int:





           '''simple docstring'''
           A_						       =   self.eval_dataset
           A_						       =   self.get_eval_dataloader(UpperCamelCase__							)

           A_						       =   next(iter(UpperCamelCase__							)							)

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

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

           model.eval()
           model.float()

           A_						       =   model.module if hasattr(UpperCamelCase__		,					"""module"""							) else model
           quant_trainer.configure_model(UpperCamelCase__		,					self.quant_trainer_args							)

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

           A_						       =   {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"""							)








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE   =  logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE   =  {
    "uw-madison/mra-base-512-4": "https://huggingface.co/uw-madison/mra-base-512-4/resolve/main/config.json",
}






class 					SCREAMING_SNAKE_CASE_  (       _a ):
       """simple docstring"""


       __lowerCAmelCase    :		str					='''mra'''





       def __init__(				self	:Any,      snake_case	:List[str]=5_0265,      snake_case	:List[Any]=768,      snake_case	:Optional[Any]=12,      snake_case	:Optional[Any]=12,      snake_case	:str=3072,      snake_case	:Tuple="gelu",      snake_case	:Optional[int]=0.1,      snake_case	:int=0.1,      snake_case	:Any=512,      snake_case	:Union[str, Any]=1,      snake_case	:Union[str, Any]=0.0_2,      snake_case	:List[Any]=1e-5,      snake_case	:Optional[int]="absolute",      snake_case	:Optional[int]=4,      snake_case	:str="full",      snake_case	:Optional[int]=0,      snake_case	:List[Any]=0,      snake_case	:int=1,      snake_case	:List[Any]=0,      snake_case	:Dict=2,      **snake_case	:Dict,      ):



          """simple docstring"""

          super().__init__(pad_token_id=snake_case,      bos_token_id=snake_case,      eos_token_id=snake_case,      **snake_case)

          _lowercase              =vocab_size
          _lowercase              =max_position_embeddings
          _lowercase              =hidden_size
          _lowercase              =num_hidden_layers
          _lowercase              =num_attention_heads
          _lowercase              =intermediate_size
          _lowercase              =hidden_act
          _lowercase              =hidden_dropout_prob
          _lowercase              =attention_probs_dropout_prob
          _lowercase              =initializer_range
          _lowercase              =type_vocab_size
          _lowercase              =layer_norm_eps
          _lowercase              =position_embedding_type
          _lowercase              =block_per_row
          _lowercase              =approx_mode
          _lowercase              =initial_prior_first_n_blocks
          _lowercase              =initial_prior_diagonal_n_blocks


