


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	 =      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 	__lowercase					(	__snake_case	):





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






         """simple docstring"""







         super().__init__(*__lowerCamelCase       ,							**__lowerCamelCase							)
         UpperCAmelCase      					=  eval_examples
         UpperCAmelCase      					=  post_process_function
         UpperCAmelCase      					=  quant_trainer_args
         UpperCAmelCase      					=  1_2_8  # default number of calibration samples





      def 							_lowercase					(		self    :       Union[str, Any]       ,							__lowerCamelCase    :       int=None							)     -> List[Any]:






         """simple docstring"""







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

         UpperCAmelCase      					=  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 							_lowercase					(		self    :       str       ,							__lowerCamelCase    :       Optional[int]=None							)     -> str:






         """simple docstring"""







         UpperCAmelCase      					=  self.train_dataset if calib_dataset is None else calib_dataset
         UpperCAmelCase      					=  self.get_calib_dataloader(__lowerCamelCase							)

         UpperCAmelCase      					=  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      					=  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      					=  model





      def 							_lowercase					(		self    :       Union[str, Any]       ,							__lowerCamelCase    :       str=None       ,							__lowerCamelCase    :       Optional[int]=None       ,							__lowerCamelCase    :       int=None       ,							__lowerCamelCase    :       str = "eval"							)     -> Optional[int]:






         """simple docstring"""







         UpperCAmelCase      					=  self.eval_dataset if eval_dataset is None else eval_dataset
         UpperCAmelCase      					=  self.get_eval_dataloader(__lowerCamelCase							)
         UpperCAmelCase      					=  self.eval_examples if eval_examples is None else eval_examples

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

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

            self.log(__lowerCamelCase							)
         else:
            UpperCAmelCase      					=  {}

         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      					=  self.callback_handler.on_evaluate(self.args       ,							self.state       ,							self.control       ,							__lowerCamelCase							)
         return metrics





      def 							_lowercase					(		self    :       Optional[Any]       ,							__lowerCamelCase    :       Optional[Any]       ,							__lowerCamelCase    :       Union[str, Any]       ,							__lowerCamelCase    :       List[str]=None       ,							__lowerCamelCase    :       str = "test"							)     -> List[Any]:






         """simple docstring"""







         UpperCAmelCase      					=  self.get_test_dataloader(__lowerCamelCase							)

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

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

         UpperCAmelCase      					=  self.post_process_function(__lowerCamelCase       ,							__lowerCamelCase       ,							output.predictions       ,							"""predict"""							)
         UpperCAmelCase      					=  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      					=  metrics.pop(__lowerCamelCase							)

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





      def 							_lowercase					(		self    :       List[str]       ,							__lowerCamelCase    :       Tuple="./"							)     -> Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase      					=  self.eval_dataset
         UpperCAmelCase      					=  self.get_eval_dataloader(__lowerCamelCase							)

         UpperCAmelCase      					=  next(iter(__lowerCamelCase							)							)

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

         # convert to tuple
         UpperCAmelCase      					=  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      					=  True

         UpperCAmelCase      					=  self.model.to(__lowerCamelCase							)

         model.eval()
         model.float()

         UpperCAmelCase      					=  model.module if hasattr(__lowerCamelCase       ,							"""module"""							) else model
         quant_trainer.configure_model(__lowerCamelCase       ,							self.quant_trainer_args							)

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

         UpperCAmelCase      					=  {0: """batch_size""", 1: """seq_len"""}

         torch.onnx.export(
             __lowerCamelCase       ,							__lowerCamelCase       ,							__lowerCamelCase       ,							export_params=__lowerCamelCase       ,							opset_version=1_3       ,							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"""							)







from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import tensorflow as tf

 from transformers import TFResNetForImageClassification, TFResNetModel
 from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor






class 	__lowercase					:





      def __init__(		self    :       str       ,							__lowerCamelCase    :       Any       ,							__lowerCamelCase    :       Optional[int]=3       ,							__lowerCamelCase    :       Union[str, Any]=3_2       ,							__lowerCamelCase    :       List[str]=3       ,							__lowerCamelCase    :       Dict=1_0       ,							__lowerCamelCase    :       Optional[int]=[1_0, 2_0, 3_0, 4_0]       ,							__lowerCamelCase    :       Optional[int]=[1, 1, 2, 1]       ,							__lowerCamelCase    :       int=True       ,							__lowerCamelCase    :       str=True       ,							__lowerCamelCase    :       str="relu"       ,							__lowerCamelCase    :       str=3       ,							__lowerCamelCase    :       int=None       ,							)     -> Dict:






         """simple docstring"""







         UpperCAmelCase      					=  parent
         UpperCAmelCase      					=  batch_size
         UpperCAmelCase      					=  image_size
         UpperCAmelCase      					=  num_channels
         UpperCAmelCase      					=  embeddings_size
         UpperCAmelCase      					=  hidden_sizes
         UpperCAmelCase      					=  depths
         UpperCAmelCase      					=  is_training
         UpperCAmelCase      					=  use_labels
         UpperCAmelCase      					=  hidden_act
         UpperCAmelCase      					=  num_labels
         UpperCAmelCase      					=  scope
         UpperCAmelCase      					=  len(__lowerCamelCase							)





      def 							_lowercase					(		self    :       List[str]							)     -> Optional[Any]:






         """simple docstring"""







         UpperCAmelCase      					=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

         UpperCAmelCase      					=  None
         if self.use_labels:
            UpperCAmelCase      					=  ids_tensor([self.batch_size]       ,							self.num_labels							)

         UpperCAmelCase      					=  self.get_config()

         return config, pixel_values, labels





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






         """simple docstring"""







         return ResNetConfig(
             num_channels=self.num_channels       ,							embeddings_size=self.embeddings_size       ,							hidden_sizes=self.hidden_sizes       ,							depths=self.depths       ,							hidden_act=self.hidden_act       ,							num_labels=self.num_labels       ,							image_size=self.image_size       ,							)





      def 							_lowercase					(		self    :       Tuple       ,							__lowerCamelCase    :       Optional[int]       ,							__lowerCamelCase    :       List[str]       ,							__lowerCamelCase    :       Tuple							)     -> str:






         """simple docstring"""







         UpperCAmelCase      					=  TFResNetModel(config=__lowerCamelCase							)
         UpperCAmelCase      					=  model(__lowerCamelCase							)
         # expected last hidden states: B, C, H // 32, W // 32
         self.parent.assertEqual(
             result.last_hidden_state.shape       ,							(self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2)       ,							)





      def 							_lowercase					(		self    :       Optional[Any]       ,							__lowerCamelCase    :       Any       ,							__lowerCamelCase    :       Optional[Any]       ,							__lowerCamelCase    :       int							)     -> Tuple:






         """simple docstring"""







         UpperCAmelCase      					=  self.num_labels
         UpperCAmelCase      					=  TFResNetForImageClassification(__lowerCamelCase							)
         UpperCAmelCase      					=  model(__lowerCamelCase       ,							labels=__lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.num_labels)							)





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






         """simple docstring"""







         UpperCAmelCase      					=  self.prepare_config_and_inputs()
         UpperCAmelCase					,	UpperCAmelCase					,	UpperCAmelCase      					=  config_and_inputs
         UpperCAmelCase      					=  {"""pixel_values""": pixel_values}
         return config, inputs_dict






@require_tf
class 	__lowercase					(	__snake_case			,       __snake_case			,       unittest.TestCase	):
      UpperCamelCase										=							(TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
      UpperCamelCase										=							(
          {'''feature-extraction''': TFResNetModel, '''image-classification''': TFResNetForImageClassification}
          if is_tf_available()
          else {}
      )

      UpperCamelCase										=							False
      UpperCamelCase										=							False
      UpperCamelCase										=							False
      UpperCamelCase										=							False
      UpperCamelCase										=							False





      def 							_lowercase					(		self    :       str							)     -> List[Any]:






         """simple docstring"""







         UpperCAmelCase      					=  TFResNetModelTester(self							)
         UpperCAmelCase      					=  ConfigTester(self       ,							config_class=__lowerCamelCase       ,							has_text_modality=__lowerCamelCase							)





      def 							_lowercase					(		self    :       str							)     -> Dict:






         """simple docstring"""







         self.create_and_test_config_common_properties()
         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()





      def 							_lowercase					(		self    :       str							)     -> Dict:






         """simple docstring"""







         return





      @unittest.skip(reason="""ResNet does not use inputs_embeds"""							)
      def 							_lowercase					(		self    :       Optional[int]							)     -> Optional[int]:






         """simple docstring"""







         pass





      @unittest.skip(reason="""ResNet does not support input and output embeddings"""							)
      def 							_lowercase					(		self    :       Union[str, Any]							)     -> Tuple:






         """simple docstring"""







         pass





      def 							_lowercase					(		self    :       Optional[Any]							)     -> Optional[int]:






         """simple docstring"""







         UpperCAmelCase					,	UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            UpperCAmelCase      					=  model_class(__lowerCamelCase							)
            UpperCAmelCase      					=  inspect.signature(model.call							)
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            UpperCAmelCase      					=  [*signature.parameters.keys()]

            UpperCAmelCase      					=  ["""pixel_values"""]
            self.assertListEqual(arg_names[:1]       ,							__lowerCamelCase							)





      def 							_lowercase					(		self    :       str							)     -> List[Any]:






         """simple docstring"""







         UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCamelCase							)





      def 							_lowercase					(		self    :       List[Any]							)     -> List[str]:






         """simple docstring"""







         def check_hidden_states_output(__lowerCamelCase    :       str       ,							__lowerCamelCase    :       Dict       ,							__lowerCamelCase    :       Union[str, Any]							):
            UpperCAmelCase      					=  model_class(__lowerCamelCase							)
            UpperCAmelCase      					=  model(**self._prepare_for_class(__lowerCamelCase       ,							__lowerCamelCase							)							)

            UpperCAmelCase      					=  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

            UpperCAmelCase      					=  self.model_tester.num_stages
            self.assertEqual(len(__lowerCamelCase							)       ,							expected_num_stages + 1							)

            # ResNet's feature maps are of shape (batch_size, num_channels, height, width)
            self.assertListEqual(
                list(hidden_states[0].shape[-2:]							)       ,							[self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,							)

         UpperCAmelCase					,	UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase      					=  ["""basic""", """bottleneck"""]
         for model_class in self.all_model_classes:
            for layer_type in layers_type:
               UpperCAmelCase      					=  layer_type
               UpperCAmelCase      					=  True
               check_hidden_states_output(__lowerCamelCase       ,							__lowerCamelCase       ,							__lowerCamelCase							)

               # check that output_hidden_states also work using config
               del inputs_dict["output_hidden_states"]
               UpperCAmelCase      					=  True

               check_hidden_states_output(__lowerCamelCase       ,							__lowerCamelCase       ,							__lowerCamelCase							)





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






         """simple docstring"""







         UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase							)





      @slow
      def 							_lowercase					(		self    :       Tuple							)     -> List[str]:






         """simple docstring"""







         for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCAmelCase      					=  TFResNetModel.from_pretrained(__lowerCamelCase							)
            self.assertIsNotNone(__lowerCamelCase							)

def 	_UpperCamelCase					(    )    ->int:
   UpperCAmelCase      					=  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
   return image







@require_tf
@require_vision
class 	__lowercase					(	unittest.TestCase	):





      @cached_property
      def 							_lowercase					(		self    :       List[Any]							)     -> Dict:






         """simple docstring"""







         return (
             AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)
             if is_vision_available()
             else None
         )





      @slow
      def 							_lowercase					(		self    :       Dict							)     -> str:






         """simple docstring"""







         UpperCAmelCase      					=  TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]							)

         UpperCAmelCase      					=  self.default_image_processor
         UpperCAmelCase      					=  prepare_img()
         UpperCAmelCase      					=  image_processor(images=__lowerCamelCase       ,							return_tensors="""tf"""							)

         # forward pass
         UpperCAmelCase      					=  model(**__lowerCamelCase							)

         # verify the logits
         UpperCAmelCase      					=  tf.TensorShape((1, 1_0_0_0)							)
         self.assertEqual(outputs.logits.shape       ,							__lowerCamelCase							)

         UpperCAmelCase      					=  tf.constant([-11.1_069, -9.7_877, -8.3_777]							)

         self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()       ,							__lowerCamelCase       ,							atol=1e-4							)							)



