


import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
 from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import ViTImageProcessor






class 	__lowercase					:





      def __init__(		self    :       Any       ,							__lowerCamelCase    :       Optional[Any]       ,							__lowerCamelCase    :       Any=1_3       ,							__lowerCamelCase    :       Optional[int]=3_0       ,							__lowerCamelCase    :       Union[str, Any]=2       ,							__lowerCamelCase    :       Union[str, Any]=3       ,							__lowerCamelCase    :       Any=True       ,							__lowerCamelCase    :       str=True       ,							__lowerCamelCase    :       Tuple=3_2       ,							__lowerCamelCase    :       int=5       ,							__lowerCamelCase    :       Optional[int]=4       ,							__lowerCamelCase    :       Optional[Any]=3_7       ,							__lowerCamelCase    :       Union[str, Any]="gelu"       ,							__lowerCamelCase    :       Any=0.1       ,							__lowerCamelCase    :       Tuple=0.1       ,							__lowerCamelCase    :       Any=1_0       ,							__lowerCamelCase    :       Tuple=0.02       ,							__lowerCamelCase    :       Dict=None       ,							__lowerCamelCase    :       Tuple=2       ,							)     -> Optional[int]:






         """simple docstring"""







         UpperCAmelCase      					=  parent
         UpperCAmelCase      					=  batch_size
         UpperCAmelCase      					=  image_size
         UpperCAmelCase      					=  patch_size
         UpperCAmelCase      					=  num_channels
         UpperCAmelCase      					=  is_training
         UpperCAmelCase      					=  use_labels
         UpperCAmelCase      					=  hidden_size
         UpperCAmelCase      					=  num_hidden_layers
         UpperCAmelCase      					=  num_attention_heads
         UpperCAmelCase      					=  intermediate_size
         UpperCAmelCase      					=  hidden_act
         UpperCAmelCase      					=  hidden_dropout_prob
         UpperCAmelCase      					=  attention_probs_dropout_prob
         UpperCAmelCase      					=  type_sequence_label_size
         UpperCAmelCase      					=  initializer_range
         UpperCAmelCase      					=  scope
         UpperCAmelCase      					=  encoder_stride

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





      def 							_lowercase					(		self    :       int							)     -> List[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.type_sequence_label_size							)

         UpperCAmelCase      					=  self.get_config()

         return config, pixel_values, labels





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






         """simple docstring"""







         return 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=__lowerCamelCase       ,							initializer_range=self.initializer_range       ,							encoder_stride=self.encoder_stride       ,							)





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






         """simple docstring"""







         UpperCAmelCase      					=  ViTModel(config=__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()
         UpperCAmelCase      					=  model(__lowerCamelCase							)
         self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)							)





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






         """simple docstring"""







         UpperCAmelCase      					=  ViTForMaskedImageModeling(config=__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()
         UpperCAmelCase      					=  model(__lowerCamelCase							)
         self.parent.assertEqual(
             result.reconstruction.shape       ,							(self.batch_size, self.num_channels, self.image_size, self.image_size)							)

         # test greyscale images
         UpperCAmelCase      					=  1
         UpperCAmelCase      					=  ViTForMaskedImageModeling(__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()

         UpperCAmelCase      					=  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
         UpperCAmelCase      					=  model(__lowerCamelCase							)
         self.parent.assertEqual(result.reconstruction.shape       ,							(self.batch_size, 1, self.image_size, self.image_size)							)





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






         """simple docstring"""







         UpperCAmelCase      					=  self.type_sequence_label_size
         UpperCAmelCase      					=  ViTForImageClassification(__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()
         UpperCAmelCase      					=  model(__lowerCamelCase       ,							labels=__lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.type_sequence_label_size)							)

         # test greyscale images
         UpperCAmelCase      					=  1
         UpperCAmelCase      					=  ViTForImageClassification(__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()

         UpperCAmelCase      					=  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
         UpperCAmelCase      					=  model(__lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.type_sequence_label_size)							)





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






         """simple docstring"""







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






@require_torch
class 	__lowercase					(	__snake_case			,       __snake_case			,       unittest.TestCase	):
      UpperCamelCase										=							(
          (
              ViTModel,
              ViTForImageClassification,
              ViTForMaskedImageModeling,
          )
          if is_torch_available()
          else ()
      )
      UpperCamelCase										=							(
          {'''feature-extraction''': ViTModel, '''image-classification''': ViTForImageClassification}
          if is_torch_available()
          else {}
      )
      UpperCamelCase										=							True

      UpperCamelCase										=							False
      UpperCamelCase										=							False
      UpperCamelCase										=							False





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






         """simple docstring"""







         UpperCAmelCase      					=  ViTModelTester(self							)
         UpperCAmelCase      					=  ConfigTester(self       ,							config_class=__lowerCamelCase       ,							has_text_modality=__lowerCamelCase       ,							hidden_size=3_7							)





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






         """simple docstring"""







         self.config_tester.run_common_tests()





      @unittest.skip(reason="""ViT does not use inputs_embeds"""							)
      def 							_lowercase					(		self    :       Tuple							)     -> Union[str, Any]:






         """simple docstring"""







         pass





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






         """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							)
            self.assertIsInstance(model.get_input_embeddings()       ,							(nn.Module)							)
            UpperCAmelCase      					=  model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(__lowerCamelCase       ,							nn.Linear							)							)





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






         """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.forward							)
            # 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    :       Optional[Any]							)     -> str:






         """simple docstring"""







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





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






         """simple docstring"""







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





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






         """simple docstring"""







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





      @slow
      def 							_lowercase					(		self    :       int							)     -> Any:






         """simple docstring"""







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

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







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





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






         """simple docstring"""







         return ViTImageProcessor.from_pretrained("""google/vit-base-patch16-224"""							) if is_vision_available() else None





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






         """simple docstring"""







         UpperCAmelCase      					=  ViTForImageClassification.from_pretrained("""google/vit-base-patch16-224"""							).to(__lowerCamelCase							)

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

         # forward pass
         with torch.no_grad():
            UpperCAmelCase      					=  model(**__lowerCamelCase							)

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

         UpperCAmelCase      					=  torch.tensor([-0.2_744, 0.8_215, -0.0_836]							).to(__lowerCamelCase							)

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





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






         """simple docstring"""







         UpperCAmelCase      					=  ViTModel.from_pretrained("""facebook/dino-vits8"""							).to(__lowerCamelCase							)

         UpperCAmelCase      					=  ViTImageProcessor.from_pretrained("""facebook/dino-vits8"""       ,							size=4_8_0							)
         UpperCAmelCase      					=  prepare_img()
         UpperCAmelCase      					=  image_processor(images=__lowerCamelCase       ,							return_tensors="""pt"""							)
         UpperCAmelCase      					=  inputs.pixel_values.to(__lowerCamelCase							)

         # forward pass
         with torch.no_grad():
            UpperCAmelCase      					=  model(__lowerCamelCase       ,							interpolate_pos_encoding=__lowerCamelCase							)

         # verify the logits
         UpperCAmelCase      					=  torch.Size((1, 3_6_0_1, 3_8_4)							)
         self.assertEqual(outputs.last_hidden_state.shape       ,							__lowerCamelCase							)

         UpperCAmelCase      					=  torch.tensor(
             [[4.2_340, 4.3_906, -6.6_692], [4.5_463, 1.8_928, -6.7_257], [4.4_429, 0.8_496, -5.8_585]]							).to(__lowerCamelCase							)

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





      @slow
      @require_accelerate
      @require_torch_gpu
      def 							_lowercase					(		self    :       List[Any]							)     -> Optional[int]:






         """simple docstring"""







         UpperCAmelCase      					=  ViTModel.from_pretrained("""facebook/dino-vits8"""       ,							torch_dtype=torch.floataa       ,							device_map="""auto"""							)
         UpperCAmelCase      					=  self.default_image_processor

         UpperCAmelCase      					=  prepare_img()
         UpperCAmelCase      					=  image_processor(images=__lowerCamelCase       ,							return_tensors="""pt"""							)
         UpperCAmelCase      					=  inputs.pixel_values.to(__lowerCamelCase							)

         # forward pass to make sure inference works in fp16
         with torch.no_grad():
            UpperCAmelCase      					=  model(__lowerCamelCase							)







import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


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






         """simple docstring"""







         super().__init__(*__lowerCamelCase       ,							**__lowerCamelCase							)
         UpperCAmelCase      					=  eval_examples
         UpperCAmelCase      					=  post_process_function





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






         """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
         UpperCAmelCase      					=  time.time()
         try:
            UpperCAmelCase      					=  eval_loop(
                __lowerCamelCase       ,							description="""Evaluation"""       ,							prediction_loss_only=True if compute_metrics is None else None       ,							ignore_keys=__lowerCamelCase       ,							metric_key_prefix=__lowerCamelCase       ,							)
         finally:
            UpperCAmelCase      					=  compute_metrics
         UpperCAmelCase      					=  self.args.eval_batch_size * self.args.world_size
         if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
            start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
         output.metrics.update(
             speed_metrics(
                 __lowerCamelCase       ,							__lowerCamelCase       ,							num_samples=output.num_samples       ,							num_steps=math.ceil(output.num_samples / total_batch_size							)       ,							)							)
         if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
            # Only the main node write the results by default
            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							)
            metrics.update(output.metrics							)
         else:
            UpperCAmelCase      					=  output.metrics

         if self.args.should_log:
            # Only the main node log the results by default
            self.log(__lowerCamelCase							)

         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    :       Union[str, Any]       ,							__lowerCamelCase    :       Any       ,							__lowerCamelCase    :       Dict=None       ,							__lowerCamelCase    :       str = "test"							)     -> Dict:






         """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
         UpperCAmelCase      					=  time.time()
         try:
            UpperCAmelCase      					=  eval_loop(
                __lowerCamelCase       ,							description="""Prediction"""       ,							prediction_loss_only=True if compute_metrics is None else None       ,							ignore_keys=__lowerCamelCase       ,							metric_key_prefix=__lowerCamelCase       ,							)
         finally:
            UpperCAmelCase      					=  compute_metrics
         UpperCAmelCase      					=  self.args.eval_batch_size * self.args.world_size
         if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
            start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
         output.metrics.update(
             speed_metrics(
                 __lowerCamelCase       ,							__lowerCamelCase       ,							num_samples=output.num_samples       ,							num_steps=math.ceil(output.num_samples / total_batch_size							)       ,							)							)

         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							)
         metrics.update(output.metrics							)
         return PredictionOutput(predictions=predictions.predictions       ,							label_ids=predictions.label_ids       ,							metrics=__lowerCamelCase							)



