
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


lowercase_ 				=			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							:							str     ,						*_A							:							Optional[int]     ,						_A							:							Optional[int]=None     ,						_A							:							Dict=None     ,						_A							:							List[Any]=None     ,						**_A							:							Optional[Any] ):







       """simple docstring"""




       super().__init__(*_A     ,						**_A )
       __SCREAMING_SNAKE_CASE      :	List[str]					=		eval_examples
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		post_process_function
       __SCREAMING_SNAKE_CASE      :	Tuple					=		quant_trainer_args
       __SCREAMING_SNAKE_CASE      :	int					=		128  # default number of calibration samples



     def  UpperCAmelCase__							(				self							:							int     ,						_A							:							Optional[int]=None ):







       """simple docstring"""




       if calib_dataset is None and self.calib_dataset is None:
         raise ValueError('''Trainer: calibration requires an calib_dataset.''' )
       __SCREAMING_SNAKE_CASE      :	Tuple					=		calib_dataset if calib_dataset is not None else self.calib_dataset

       __SCREAMING_SNAKE_CASE      :	str					=		self._remove_unused_columns(_A     ,						description='''Calibration''' )

       return DataLoader(
           _A     ,						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=_A     ,						)



     def  UpperCAmelCase__							(				self							:							Optional[int]     ,						_A							:							Tuple=None ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[str]					=		self.train_dataset if calib_dataset is None else calib_dataset
       __SCREAMING_SNAKE_CASE      :	Tuple					=		self.get_calib_dataloader(_A )

       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		self.model
       quant_trainer.configure_model(_A     ,						self.quant_trainer_args     ,						calib=_A )
       model.eval()
       quant_trainer.enable_calibration(_A )

       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(_A ):
         # Prediction step
         __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE      :	List[str]					=		self.prediction_step(_A     ,						_A     ,						prediction_loss_only=_A )
         if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
           break

       quant_trainer.finish_calibration(_A     ,						self.quant_trainer_args )
       __SCREAMING_SNAKE_CASE      :	Dict					=		model



     def  UpperCAmelCase__							(				self							:							List[str]     ,						_A							:							int=None     ,						_A							:							Optional[Any]=None     ,						_A							:							List[Any]=None     ,						_A							:							str = "eval" ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[str]					=		self.eval_dataset if eval_dataset is None else eval_dataset
       __SCREAMING_SNAKE_CASE      :	str					=		self.get_eval_dataloader(_A )
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		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      :	List[str]					=		self.compute_metrics
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		None
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
       try:
         __SCREAMING_SNAKE_CASE      :	Optional[int]					=		eval_loop(
             _A     ,						description='''Evaluation'''     ,						prediction_loss_only=True if compute_metrics is None else None     ,						ignore_keys=_A     ,						)
       finally:
         __SCREAMING_SNAKE_CASE      :	List[str]					=		compute_metrics

       if self.post_process_function is not None and self.compute_metrics is not None:
         __SCREAMING_SNAKE_CASE      :	Dict					=		self.post_process_function(_A     ,						_A     ,						output.predictions )
         __SCREAMING_SNAKE_CASE      :	str					=		self.compute_metrics(_A )

         # Prefix all keys with metric_key_prefix + '_'
         for key in list(metrics.keys() ):
           if not key.startswith(F'''{metric_key_prefix}_''' ):
             __SCREAMING_SNAKE_CASE      :	Tuple					=		metrics.pop(_A )

         self.log(_A )
       else:
         __SCREAMING_SNAKE_CASE      :	str					=		{}

       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      :	List[str]					=		self.callback_handler.on_evaluate(self.args     ,						self.state     ,						self.control     ,						_A )
       return metrics



     def  UpperCAmelCase__							(				self							:							Tuple     ,						_A							:							int     ,						_A							:							Optional[int]     ,						_A							:							str=None     ,						_A							:							str = "test" ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	str					=		self.get_test_dataloader(_A )

       # Temporarily disable metric computation, we will do it in the loop here.
       __SCREAMING_SNAKE_CASE      :	str					=		self.compute_metrics
       __SCREAMING_SNAKE_CASE      :	Dict					=		None
       __SCREAMING_SNAKE_CASE      :	Tuple					=		self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
       try:
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		eval_loop(
             _A     ,						description='''Prediction'''     ,						prediction_loss_only=True if compute_metrics is None else None     ,						ignore_keys=_A     ,						)
       finally:
         __SCREAMING_SNAKE_CASE      :	Tuple					=		compute_metrics

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

       __SCREAMING_SNAKE_CASE      :	Dict					=		self.post_process_function(_A     ,						_A     ,						output.predictions     ,						'''predict''' )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.compute_metrics(_A )

       # Prefix all keys with metric_key_prefix + '_'
       for key in list(metrics.keys() ):
         if not key.startswith(F'''{metric_key_prefix}_''' ):
           __SCREAMING_SNAKE_CASE      :	List[str]					=		metrics.pop(_A )

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




     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							List[Any]="./" ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		self.eval_dataset
       __SCREAMING_SNAKE_CASE      :	int					=		self.get_eval_dataloader(_A )

       __SCREAMING_SNAKE_CASE      :	str					=		next(iter(_A ) )

       # saving device - to make it consistent
       __SCREAMING_SNAKE_CASE      :	str					=		torch.device('''cuda''' if torch.cuda.is_available() else '''cpu''' )

       # convert to tuple
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		tuple(v.to(_A ) for k, v in batch.items() )

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

       __SCREAMING_SNAKE_CASE      :	Dict					=		True

       __SCREAMING_SNAKE_CASE      :	Dict					=		self.model.to(_A )

       model.eval()
       model.float()

       __SCREAMING_SNAKE_CASE      :	str					=		model.module if hasattr(_A     ,						'''module''' ) else model
       quant_trainer.configure_model(_A     ,						self.quant_trainer_args )

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		os.path.join(_A     ,						'''model.onnx''' )
       logger.info(F'''exporting model to {output_model_file}''' )

       __SCREAMING_SNAKE_CASE      :	List[str]					=		{0: '''batch_size''', 1: '''seq_len'''}

       torch.onnx.export(
           _A     ,						_A     ,						_A     ,						export_params=_A     ,						opset_version=13     ,						do_constant_folding=_A     ,						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=_A     ,						)
       logger.info('''onnx export finished''' )



"""simple docstring"""




import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
						import jax

						from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel





class  lowerCamelCase    (unittest.TestCase      ):

			def __init__(       self					:    List[Any]     ,					__UpperCAmelCase					:    Dict     ,					__UpperCAmelCase					:    Optional[Any]=1_3     ,					__UpperCAmelCase					:    Tuple=3_0     ,					__UpperCAmelCase					:    str=2     ,					__UpperCAmelCase					:    List[Any]=3     ,					__UpperCAmelCase					:    Any=True     ,					__UpperCAmelCase					:    Dict=True     ,					__UpperCAmelCase					:    Optional[Any]=3_2     ,					__UpperCAmelCase					:    Any=5     ,					__UpperCAmelCase					:    Optional[Any]=4     ,					__UpperCAmelCase					:    Tuple=3_7     ,					__UpperCAmelCase					:    List[str]="gelu"     ,					__UpperCAmelCase					:    Union[str, Any]=0.1     ,					__UpperCAmelCase					:    List[str]=0.1     ,					__UpperCAmelCase					:    int=1_0     ,					__UpperCAmelCase					:    Any=0.02     ,					) ->     Union[str, Any]:
						SCREAMING_SNAKE_CASE__										=   parent
						SCREAMING_SNAKE_CASE__										=   batch_size
						SCREAMING_SNAKE_CASE__										=   image_size
						SCREAMING_SNAKE_CASE__										=   patch_size
						SCREAMING_SNAKE_CASE__										=   num_channels
						SCREAMING_SNAKE_CASE__										=   is_training
						SCREAMING_SNAKE_CASE__										=   use_labels
						SCREAMING_SNAKE_CASE__										=   hidden_size
						SCREAMING_SNAKE_CASE__										=   num_hidden_layers
						SCREAMING_SNAKE_CASE__										=   num_attention_heads
						SCREAMING_SNAKE_CASE__										=   intermediate_size
						SCREAMING_SNAKE_CASE__										=   hidden_act
						SCREAMING_SNAKE_CASE__										=   hidden_dropout_prob
						SCREAMING_SNAKE_CASE__										=   attention_probs_dropout_prob
						SCREAMING_SNAKE_CASE__										=   type_sequence_label_size
						SCREAMING_SNAKE_CASE__										=   initializer_range

						# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
						SCREAMING_SNAKE_CASE__										=   (image_size // patch_size) ** 2
						SCREAMING_SNAKE_CASE__										=   num_patches + 1

			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     Any:
						SCREAMING_SNAKE_CASE__										=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

						SCREAMING_SNAKE_CASE__										=   ViTConfig(
						    image_size=self.image_size     ,					patch_size=self.patch_size     ,					num_channels=self.num_channels     ,					hidden_size=self.hidden_size     ,					num_hidden_layers=self.num_hidden_layers     ,					num_attention_heads=self.num_attention_heads     ,					intermediate_size=self.intermediate_size     ,					hidden_act=self.hidden_act     ,					hidden_dropout_prob=self.hidden_dropout_prob     ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,					is_decoder=__UpperCAmelCase     ,					initializer_range=self.initializer_range     ,					)

						return config, pixel_values

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict     ,					__UpperCAmelCase					:    int     ,					__UpperCAmelCase					:    int							) ->     Optional[Any]:
						SCREAMING_SNAKE_CASE__										=   FlaxViTModel(config=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   model(__UpperCAmelCase							)
						# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
						SCREAMING_SNAKE_CASE__										=   (self.image_size, self.image_size)
						SCREAMING_SNAKE_CASE__										=   (self.patch_size, self.patch_size)
						SCREAMING_SNAKE_CASE__										=   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
						self.parent.assertEqual(result.last_hidden_state.shape     ,					(self.batch_size, num_patches + 1, self.hidden_size)							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]     ,					__UpperCAmelCase					:    Optional[int]     ,					__UpperCAmelCase					:    Optional[Any]							) ->     int:
						SCREAMING_SNAKE_CASE__										=   self.type_sequence_label_size
						SCREAMING_SNAKE_CASE__										=   FlaxViTForImageClassification(config=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   model(__UpperCAmelCase							)
						self.parent.assertEqual(result.logits.shape     ,					(self.batch_size, self.type_sequence_label_size)							)

						# test greyscale images
						SCREAMING_SNAKE_CASE__										=   1
						SCREAMING_SNAKE_CASE__										=   FlaxViTForImageClassification(__UpperCAmelCase							)

						SCREAMING_SNAKE_CASE__										=   floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
						SCREAMING_SNAKE_CASE__										=   model(__UpperCAmelCase							)





			def        SCREAMING_SNAKE_CASE					(       self					:    Tuple							) ->     str:
						SCREAMING_SNAKE_CASE__										=   self.prepare_config_and_inputs()
						(
						    (
						    SCREAMING_SNAKE_CASE__
						)				,		(
						    SCREAMING_SNAKE_CASE__
						)				,		
						)										=   config_and_inputs
						SCREAMING_SNAKE_CASE__										=   {"""pixel_values""": pixel_values}
						return config, inputs_dict





@require_flax
class  lowerCamelCase    (A__      ,unittest.TestCase      ):
			lowerCamelCase__ :		Tuple       		=				(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()

			def        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     None:
						SCREAMING_SNAKE_CASE__										=   FlaxViTModelTester(self							)
						SCREAMING_SNAKE_CASE__										=   ConfigTester(self     ,					config_class=__UpperCAmelCase     ,					has_text_modality=__UpperCAmelCase     ,					hidden_size=3_7							)

			def        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     Dict:
						self.config_tester.run_common_tests()

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]							) ->     Optional[Any]:
						SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     List[str]:
						SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     str:
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs_for_common()

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

									SCREAMING_SNAKE_CASE__										=   ["""pixel_values"""]
									self.assertListEqual(arg_names[:1]     ,					__UpperCAmelCase							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     Optional[Any]:
						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__							):
												SCREAMING_SNAKE_CASE__										=   self._prepare_for_class(__UpperCAmelCase     ,					__UpperCAmelCase							)
												SCREAMING_SNAKE_CASE__										=   model_class(__UpperCAmelCase							)

												@jax.jit
												def model_jitted(__UpperCAmelCase					:    int     ,					**__UpperCAmelCase					:    Tuple							):
															return model(pixel_values=__UpperCAmelCase     ,					**__UpperCAmelCase							)

												with self.subTest("""JIT Enabled"""							):
															SCREAMING_SNAKE_CASE__										=   model_jitted(**__UpperCAmelCase							).to_tuple()

												with self.subTest("""JIT Disabled"""							):
															with jax.disable_jit():
																		SCREAMING_SNAKE_CASE__										=   model_jitted(**__UpperCAmelCase							).to_tuple()

												self.assertEqual(len(__UpperCAmelCase							)     ,					len(__UpperCAmelCase							)							)
												for jitted_output, output in zip(__UpperCAmelCase     ,					__UpperCAmelCase							):
															self.assertEqual(jitted_output.shape     ,					output.shape							)





			@slow
			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     Dict:
						for model_class_name in self.all_model_classes:
									SCREAMING_SNAKE_CASE__										=   model_class_name.from_pretrained("""google/vit-base-patch16-224"""							)
									SCREAMING_SNAKE_CASE__										=   model(np.ones((1, 3, 2_2_4, 2_2_4)							)							)
									self.assertIsNotNone(__UpperCAmelCase							)


