





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


__magic_name__   :						Dict    =			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			):



     '''simple docstring'''






     def __init__(					self						,   *lowerCamelCase						,   lowerCamelCase=None						,   lowerCamelCase=None						,   lowerCamelCase=None						,   **lowerCamelCase	):
       super().__init__(*lowerCamelCase						,   **lowerCamelCase	)
       _snake_case     		=				eval_examples
       _snake_case     		=				post_process_function
       _snake_case     		=				quant_trainer_args
       _snake_case     		=				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."	)
       _snake_case     		=				calib_dataset if calib_dataset is not None else self.calib_dataset

       _snake_case     		=				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	):
       _snake_case     		=				self.train_dataset if calib_dataset is None else calib_dataset
       _snake_case     		=				self.get_calib_dataloader(lowerCamelCase	)

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

     def 	UpperCamelCase(					self						,   lowerCamelCase=None						,   lowerCamelCase=None						,   lowerCamelCase=None						,   lowerCamelCase = "eval"	):
       _snake_case     		=				self.eval_dataset if eval_dataset is None else eval_dataset
       _snake_case     		=				self.get_eval_dataloader(lowerCamelCase	)
       _snake_case     		=				self.eval_examples if eval_examples is None else eval_examples

       # Temporarily disable metric computation, we will do it in the loop here.
       _snake_case     		=				self.compute_metrics
       _snake_case     		=				None
       _snake_case     		=				self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
       try:
         _snake_case     		=				eval_loop(
             lowerCamelCase						,   description="Evaluation"						,   prediction_loss_only=True if compute_metrics is None else None						,   ignore_keys=lowerCamelCase						,   )
       finally:
         _snake_case     		=				compute_metrics

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

         self.log(lowerCamelCase	)
       else:
         _snake_case     		=				{}

       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()	)

       _snake_case     		=				self.callback_handler.on_evaluate(self.args						,   self.state						,   self.control						,   lowerCamelCase	)
       return metrics

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase=None						,   lowerCamelCase = "test"	):
       _snake_case     		=				self.get_test_dataloader(lowerCamelCase	)

       # Temporarily disable metric computation, we will do it in the loop here.
       _snake_case     		=				self.compute_metrics
       _snake_case     		=				None
       _snake_case     		=				self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
       try:
         _snake_case     		=				eval_loop(
             lowerCamelCase						,   description="Prediction"						,   prediction_loss_only=True if compute_metrics is None else None						,   ignore_keys=lowerCamelCase						,   )
       finally:
         _snake_case     		=				compute_metrics

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

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

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






     def 	UpperCamelCase(					self						,   lowerCamelCase="./"	):
       _snake_case     		=				self.eval_dataset
       _snake_case     		=				self.get_eval_dataloader(lowerCamelCase	)

       _snake_case     		=				next(iter(lowerCamelCase	)	)

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

       # convert to tuple
       _snake_case     		=				tuple(v.to(lowerCamelCase	) for k, v in batch.items()	)

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

       _snake_case     		=				True

       _snake_case     		=				self.model.to(lowerCamelCase	)

       model.eval()
       model.float()

       _snake_case     		=				model.module if hasattr(lowerCamelCase						,   "module"	) else model
       quant_trainer.configure_model(lowerCamelCase						,   self.quant_trainer_args	)

       _snake_case     		=				os.path.join(lowerCamelCase						,   "model.onnx"	)
       logger.info(F'''exporting model to {output_model_file}'''	)

       _snake_case     		=				{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'''



from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


__magic_name__   :						Tuple    =			logging.get_logger(__name__)


if is_vision_available():
 import PIL


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						str      =   ['''pixel_values''']

     def __init__(					self						,   lowerCamelCase = True						,   lowerCamelCase = None						,   lowerCamelCase = PILImageResampling.BICUBIC						,   lowerCamelCase = True						,   lowerCamelCase = None						,   lowerCamelCase = True						,   lowerCamelCase = 1 / 255						,   lowerCamelCase = True						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = True						,   **lowerCamelCase						,   ):
       super().__init__(**lowerCamelCase	)
       _snake_case     		=				size if size is not None else {"shortest_edge": 224}
       _snake_case     		=				get_size_dict(lowerCamelCase						,   default_to_square=lowerCamelCase	)
       _snake_case     		=				crop_size if crop_size is not None else {"height": 224, "width": 224}
       _snake_case     		=				get_size_dict(lowerCamelCase						,   default_to_square=lowerCamelCase						,   param_name="crop_size"	)

       _snake_case     		=				do_resize
       _snake_case     		=				size
       _snake_case     		=				resample
       _snake_case     		=				do_center_crop
       _snake_case     		=				crop_size
       _snake_case     		=				do_rescale
       _snake_case     		=				rescale_factor
       _snake_case     		=				do_normalize
       _snake_case     		=				image_mean if image_mean is not None else OPENAI_CLIP_MEAN
       _snake_case     		=				image_std if image_std is not None else OPENAI_CLIP_STD
       _snake_case     		=				do_convert_rgb

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase = PILImageResampling.BICUBIC						,   lowerCamelCase = None						,   **lowerCamelCase						,   ):
       _snake_case     		=				get_size_dict(lowerCamelCase						,   default_to_square=lowerCamelCase	)
       if "shortest_edge" not in size:
         raise ValueError(F'''The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'''	)
       _snake_case     		=				get_resize_output_image_size(lowerCamelCase						,   size=size["shortest_edge"]						,   default_to_square=lowerCamelCase	)
       return resize(lowerCamelCase						,   size=lowerCamelCase						,   resample=lowerCamelCase						,   data_format=lowerCamelCase						,   **lowerCamelCase	)

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase = None						,   **lowerCamelCase						,   ):
       _snake_case     		=				get_size_dict(lowerCamelCase	)
       if "height" not in size or "width" not in size:
         raise ValueError(F'''The `size` parameter must contain the keys (height, width). Got {size.keys()}'''	)
       return center_crop(lowerCamelCase						,   size=(size["height"], size["width"])						,   data_format=lowerCamelCase						,   **lowerCamelCase	)

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase = None						,   **lowerCamelCase						,   ):
       return rescale(lowerCamelCase						,   scale=lowerCamelCase						,   data_format=lowerCamelCase						,   **lowerCamelCase	)

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase = None						,   **lowerCamelCase						,   ):
       return normalize(lowerCamelCase						,   mean=lowerCamelCase						,   std=lowerCamelCase						,   data_format=lowerCamelCase						,   **lowerCamelCase	)






     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = None						,   lowerCamelCase = ChannelDimension.FIRST						,   **lowerCamelCase						,   ):
       _snake_case     		=				do_resize if do_resize is not None else self.do_resize
       _snake_case     		=				size if size is not None else self.size
       _snake_case     		=				get_size_dict(lowerCamelCase						,   param_name="size"						,   default_to_square=lowerCamelCase	)
       _snake_case     		=				resample if resample is not None else self.resample
       _snake_case     		=				do_center_crop if do_center_crop is not None else self.do_center_crop
       _snake_case     		=				crop_size if crop_size is not None else self.crop_size
       _snake_case     		=				get_size_dict(lowerCamelCase						,   param_name="crop_size"						,   default_to_square=lowerCamelCase	)
       _snake_case     		=				do_rescale if do_rescale is not None else self.do_rescale
       _snake_case     		=				rescale_factor if rescale_factor is not None else self.rescale_factor
       _snake_case     		=				do_normalize if do_normalize is not None else self.do_normalize
       _snake_case     		=				image_mean if image_mean is not None else self.image_mean
       _snake_case     		=				image_std if image_std is not None else self.image_std
       _snake_case     		=				do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

       _snake_case     		=				make_list_of_images(lowerCamelCase	)

       if not valid_images(lowerCamelCase	):
         raise ValueError(
             "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
             "torch.Tensor, tf.Tensor or jax.ndarray."	)

       if do_resize and size is None:
         raise ValueError("Size must be specified if do_resize is True."	)

       if do_center_crop and crop_size is None:
         raise ValueError("Crop size must be specified if do_center_crop is True."	)

       if do_rescale and rescale_factor is None:
         raise ValueError("Rescale factor must be specified if do_rescale is True."	)

       if do_normalize and (image_mean is None or image_std is None):
         raise ValueError("Image mean and std must be specified if do_normalize is True."	)

       # PIL RGBA images are converted to RGB
       if do_convert_rgb:
         _snake_case     		=				[convert_to_rgb(lowerCamelCase	) for image in images]

       # All transformations expect numpy arrays.
       _snake_case     		=				[to_numpy_array(lowerCamelCase	) for image in images]

       if do_resize:
         _snake_case     		=				[self.resize(image=lowerCamelCase						,   size=lowerCamelCase						,   resample=lowerCamelCase	) for image in images]

       if do_center_crop:
         _snake_case     		=				[self.center_crop(image=lowerCamelCase						,   size=lowerCamelCase	) for image in images]

       if do_rescale:
         _snake_case     		=				[self.rescale(image=lowerCamelCase						,   scale=lowerCamelCase	) for image in images]

       if do_normalize:
         _snake_case     		=				[self.normalize(image=lowerCamelCase						,   mean=lowerCamelCase						,   std=lowerCamelCase	) for image in images]

       _snake_case     		=				[to_channel_dimension_format(lowerCamelCase						,   lowerCamelCase	) for image in images]

       _snake_case     		=				{"pixel_values": images}
       return BatchFeature(data=lowerCamelCase						,   tensor_type=lowerCamelCase	)


