





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


SCREAMING_SNAKE_CASE_						       =	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 		UpperCamelCase__			(				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   Optional[Any]						,*lowerCamelCase__						:   int						,lowerCamelCase__						:   str=None						,lowerCamelCase__						:   List[Any]=None						,lowerCamelCase__						:   Tuple=None						,**lowerCamelCase__						:   str	)					->		int:




           '''simple docstring'''


           super().__init__(*lowerCamelCase__						,**lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						eval_examples
           SCREAMING_SNAKE_CASE		=						post_process_function
           SCREAMING_SNAKE_CASE		=						quant_trainer_args
           SCREAMING_SNAKE_CASE		=						128  # default number of calibration samples






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]						,lowerCamelCase__						:   Optional[Any]=None	)					->		int:




           '''simple docstring'''


           if calib_dataset is None and self.calib_dataset is None:
                 raise ValueError("""Trainer: calibration requires an calib_dataset."""	)
           SCREAMING_SNAKE_CASE		=						calib_dataset if calib_dataset is not None else self.calib_dataset

           SCREAMING_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     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]						,lowerCamelCase__						:   List[str]=None	)					->		Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.train_dataset if calib_dataset is None else calib_dataset
           SCREAMING_SNAKE_CASE		=						self.get_calib_dataloader(lowerCamelCase__	)

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






     def     SCREAMING_SNAKE_CASE__  (							self						:   int						,lowerCamelCase__						:   List[Any]=None						,lowerCamelCase__						:   Union[str, Any]=None						,lowerCamelCase__						:   Any=None						,lowerCamelCase__						:   str = "eval"	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.eval_dataset if eval_dataset is None else eval_dataset
           SCREAMING_SNAKE_CASE		=						self.get_eval_dataloader(lowerCamelCase__	)
           SCREAMING_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.
           SCREAMING_SNAKE_CASE		=						self.compute_metrics
           SCREAMING_SNAKE_CASE		=						None
           SCREAMING_SNAKE_CASE		=						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
           try:
                 SCREAMING_SNAKE_CASE		=						eval_loop(
                     lowerCamelCase__						,description="""Evaluation"""						,prediction_loss_only=True if compute_metrics is None else None						,ignore_keys=lowerCamelCase__						,)
           finally:
                 SCREAMING_SNAKE_CASE		=						compute_metrics

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

                 self.log(lowerCamelCase__	)
           else:
                 SCREAMING_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()	)

           SCREAMING_SNAKE_CASE		=						self.callback_handler.on_evaluate(self.args						,self.state						,self.control						,lowerCamelCase__	)
           return metrics






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]						,lowerCamelCase__						:   Optional[int]						,lowerCamelCase__						:   Optional[int]						,lowerCamelCase__						:   Optional[Any]=None						,lowerCamelCase__						:   str = "test"	)					->		Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.get_test_dataloader(lowerCamelCase__	)

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

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

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

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






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]						,lowerCamelCase__						:   Dict="./"	)					->		int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.eval_dataset
           SCREAMING_SNAKE_CASE		=						self.get_eval_dataloader(lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						next(iter(lowerCamelCase__	)	)

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

           # convert to tuple
           SCREAMING_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

           SCREAMING_SNAKE_CASE		=						True

           SCREAMING_SNAKE_CASE		=						self.model.to(lowerCamelCase__	)

           model.eval()
           model.float()

           SCREAMING_SNAKE_CASE		=						model.module if hasattr(lowerCamelCase__						,"""module"""	) else model
           quant_trainer.configure_model(lowerCamelCase__						,self.quant_trainer_args	)

           SCREAMING_SNAKE_CASE		=						os.path.join(lowerCamelCase__						,"""model.onnx"""	)
           logger.info(F"""exporting model to {output_model_file}"""	)

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









import unittest

import numpy as np

from transformers.file_utils import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import DPTImageProcessor

class 		UpperCamelCase__			(				unittest.TestCase       ):



     '''simple docstring'''

     def __init__(							self						:   Optional[Any]						,lowerCamelCase__						:   List[str]						,lowerCamelCase__						:   Optional[Any]=7						,lowerCamelCase__						:   Optional[Any]=3						,lowerCamelCase__						:   Union[str, Any]=18						,lowerCamelCase__						:   Optional[int]=30						,lowerCamelCase__						:   int=400						,lowerCamelCase__						:   List[str]=True						,lowerCamelCase__						:   Union[str, Any]=None						,lowerCamelCase__						:   Union[str, Any]=True						,lowerCamelCase__						:   Optional[int]=[0.5, 0.5, 0.5]						,lowerCamelCase__						:   Any=[0.5, 0.5, 0.5]						,)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						size if size is not None else {"""height""": 18, """width""": 18}
           SCREAMING_SNAKE_CASE		=						parent
           SCREAMING_SNAKE_CASE		=						batch_size
           SCREAMING_SNAKE_CASE		=						num_channels
           SCREAMING_SNAKE_CASE		=						image_size
           SCREAMING_SNAKE_CASE		=						min_resolution
           SCREAMING_SNAKE_CASE		=						max_resolution
           SCREAMING_SNAKE_CASE		=						do_resize
           SCREAMING_SNAKE_CASE		=						size
           SCREAMING_SNAKE_CASE		=						do_normalize
           SCREAMING_SNAKE_CASE		=						image_mean
           SCREAMING_SNAKE_CASE		=						image_std






     def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		Union[str, Any]:




           '''simple docstring'''


           return {
               "image_mean": self.image_mean,
               "image_std": self.image_std,
               "do_normalize": self.do_normalize,
               "do_resize": self.do_resize,
               "size": self.size,
           }

@require_torch
@require_vision
class 		UpperCamelCase__			(				lowerCAmelCase_							,				unittest.TestCase       ):



     '''simple docstring'''

     __snake_case   :     Tuple		     =      DPTImageProcessor if is_vision_available() else None






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						DPTImageProcessingTester(self	)






     @property
     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		int:




           '''simple docstring'''


           return self.image_processor_tester.prepare_image_processor_dict()






     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""image_mean"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""image_std"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""do_normalize"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""do_resize"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""size"""	)	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   int	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class.from_dict(self.image_processor_dict	)
           self.assertEqual(image_processor.size						,{"""height""": 18, """width""": 18}	)

           SCREAMING_SNAKE_CASE		=						self.image_processing_class.from_dict(self.image_processor_dict						,size=42	)
           self.assertEqual(image_processor.size						,{"""height""": 42, """width""": 42}	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           # create random PIL images
           SCREAMING_SNAKE_CASE		=						prepare_image_inputs(self.image_processor_tester						,equal_resolution=lowerCamelCase__	)
           for image in image_inputs:
                 self.assertIsInstance(lowerCamelCase__						,Image.Image	)

           # Test not batched input
           SCREAMING_SNAKE_CASE		=						image_processing(image_inputs[0]						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)

           # Test batched
           SCREAMING_SNAKE_CASE		=						image_processing(lowerCamelCase__						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           # create random numpy tensors
           SCREAMING_SNAKE_CASE		=						prepare_image_inputs(self.image_processor_tester						,equal_resolution=lowerCamelCase__						,numpify=lowerCamelCase__	)
           for image in image_inputs:
                 self.assertIsInstance(lowerCamelCase__						,np.ndarray	)

           # Test not batched input
           SCREAMING_SNAKE_CASE		=						image_processing(image_inputs[0]						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)

           # Test batched
           SCREAMING_SNAKE_CASE		=						image_processing(lowerCamelCase__						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           # create random PyTorch tensors
           SCREAMING_SNAKE_CASE		=						prepare_image_inputs(self.image_processor_tester						,equal_resolution=lowerCamelCase__						,torchify=lowerCamelCase__	)
           for image in image_inputs:
                 self.assertIsInstance(lowerCamelCase__						,torch.Tensor	)

           # Test not batched input
           SCREAMING_SNAKE_CASE		=						image_processing(image_inputs[0]						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)

           # Test batched
           SCREAMING_SNAKE_CASE		=						image_processing(lowerCamelCase__						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)


