

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   UpperCamelCase			:

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

     """simple docstring"""






     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
     SCREAMING_SNAKE_CASE		    =      scope
     SCREAMING_SNAKE_CASE		    =      encoder_stride

     # 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 UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE		    =      None
     if self.use_labels:
      SCREAMING_SNAKE_CASE		    =      ids_tensor([self.batch_size]				,     self.type_sequence_label_size     )

     SCREAMING_SNAKE_CASE		    =      self.get_config()

     return config, pixel_values, labels

    def UpperCamelCase      (      self			:   int     ):

     """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=snake_case__				,     initializer_range=self.initializer_range				,     encoder_stride=self.encoder_stride				,     )

    def UpperCamelCase      (      self			:   Optional[int]				,     snake_case__			:   str				,     snake_case__			:   Tuple				,     snake_case__			:   Tuple     ):

     """simple docstring"""






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

    def UpperCamelCase      (      self			:   str				,     snake_case__			:   Union[str, Any]				,     snake_case__			:   Optional[Any]				,     snake_case__			:   Tuple     ):

     """simple docstring"""






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

     # test greyscale images
     SCREAMING_SNAKE_CASE		    =      1
     SCREAMING_SNAKE_CASE		    =      ViTForMaskedImageModeling(snake_case__     )
     model.to(snake_case__     )
     model.eval()

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

    def UpperCamelCase      (      self			:   str				,     snake_case__			:   Any				,     snake_case__			:   int				,     snake_case__			:   Union[str, Any]     ):

     """simple docstring"""






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

     # test greyscale images
     SCREAMING_SNAKE_CASE		    =      1
     SCREAMING_SNAKE_CASE		    =      ViTForImageClassification(snake_case__     )
     model.to(snake_case__     )
     model.eval()

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




    def UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.prepare_config_and_inputs()
     (
         (
         SCREAMING_SNAKE_CASE
     )    ,						(
         SCREAMING_SNAKE_CASE
     )    ,						(
         SCREAMING_SNAKE_CASE
     )    ,						
     )		    =      config_and_inputs
     SCREAMING_SNAKE_CASE		    =      {'pixel_values': pixel_values}
     return config, inputs_dict




@require_torch
class   UpperCamelCase			(  SCREAMING_SNAKE_CASE						,      SCREAMING_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 UpperCamelCase      (      self			:   Union[str, Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      ViTModelTester(self     )
     SCREAMING_SNAKE_CASE		    =      ConfigTester(self				,     config_class=snake_case__				,     has_text_modality=snake_case__				,     hidden_size=3_7     )

    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






     self.config_tester.run_common_tests()

    @unittest.skip(reason='ViT does not use inputs_embeds'     )
    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     pass

    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






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

    def UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






     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(snake_case__     )
      SCREAMING_SNAKE_CASE		    =      inspect.signature(model.forward     )
      # 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]				,     snake_case__     )

    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*snake_case__     )

    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*snake_case__     )

    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*snake_case__     )




    @slow
    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






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


def 							__lowerCAmelCase     (							)		->				Dict:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
 return image




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

    @cached_property
    def UpperCamelCase      (      self			:   Optional[Any]     ):

     """simple docstring"""






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

    @slow
    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'     ).to(snake_case__     )

     SCREAMING_SNAKE_CASE		    =      self.default_image_processor
     SCREAMING_SNAKE_CASE		    =      prepare_img()
     SCREAMING_SNAKE_CASE		    =      image_processor(images=snake_case__				,     return_tensors='pt'     ).to(snake_case__     )

     # forward pass
     with torch.no_grad():
      SCREAMING_SNAKE_CASE		    =      model(**snake_case__     )

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

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

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

    @slow
    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      ViTModel.from_pretrained('facebook/dino-vits8'     ).to(snake_case__     )

     SCREAMING_SNAKE_CASE		    =      ViTImageProcessor.from_pretrained('facebook/dino-vits8'				,     size=4_8_0     )
     SCREAMING_SNAKE_CASE		    =      prepare_img()
     SCREAMING_SNAKE_CASE		    =      image_processor(images=snake_case__				,     return_tensors='pt'     )
     SCREAMING_SNAKE_CASE		    =      inputs.pixel_values.to(snake_case__     )

     # forward pass
     with torch.no_grad():
      SCREAMING_SNAKE_CASE		    =      model(snake_case__				,     interpolate_pos_encoding=snake_case__     )

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

     SCREAMING_SNAKE_CASE		    =      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(snake_case__     )

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




    @slow
    @require_accelerate
    @require_torch_gpu
    def UpperCamelCase      (      self			:   List[str]     ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE		    =      prepare_img()
     SCREAMING_SNAKE_CASE		    =      image_processor(images=snake_case__				,     return_tensors='pt'     )
     SCREAMING_SNAKE_CASE		    =      inputs.pixel_values.to(snake_case__     )

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




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


a_     :    str     		=   [
    "EAGER",
    "AOT_EAGER",
    "INDUCTOR",
    "NVFUSER",
    "AOT_NVFUSER",
    "AOT_CUDAGRAPHS",
    "OFI",
    "FX2TRT",
    "ONNXRT",
    "IPEX",
]


def 							__lowerCAmelCase     (							_UpperCamelCase		:	int		,				_UpperCamelCase		:	List[Any]=None		,				_UpperCamelCase		:	Dict=None		,				_UpperCamelCase		:	str=None						)		->				Any:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      True
 while ask_again:
  SCREAMING_SNAKE_CASE		    =      input(_UpperCamelCase						)
  try:
   if default is not None and len(_UpperCamelCase						) == 0:
    return default
   return convert_value(_UpperCamelCase						) if convert_value is not None else result
  except Exception:
   if error_message is not None:
    print(_UpperCamelCase						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Optional[Any]		,				_UpperCamelCase		:	Union[str, Any]=[]		,				_UpperCamelCase		:	Optional[int]=None		,				_UpperCamelCase		:	Optional[Any]=0						)		->				List[Any]:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      BulletMenu(_UpperCamelCase		,				_UpperCamelCase						)
 SCREAMING_SNAKE_CASE		    =      menu.run(default_choice=_UpperCamelCase						)
 return convert_value(_UpperCamelCase						) if convert_value is not None else result


def 							__lowerCAmelCase     (							_UpperCamelCase		:	List[str]						)		->				int:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return ComputeEnvironment(['LOCAL_MACHINE', 'AMAZON_SAGEMAKER'][value]						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Tuple						)		->				Tuple:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return DistributedType(['NO', 'MULTI_CPU', 'MULTI_XPU', 'MULTI_GPU', 'MULTI_NPU', 'TPU'][value]						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	List[Any]						)		->				List[str]:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return DynamoBackend(DYNAMO_BACKENDS[value]						).value


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Union[str, Any]						)		->				Dict:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return PrecisionType(['no', 'fp16', 'bf16', 'fp8'][value]						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Tuple						)		->				str:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE		    =      int(_UpperCamelCase						)
 return SageMakerDistributedType(['NO', 'DATA_PARALLEL', 'MODEL_PARALLEL'][value]						)


def 							__lowerCAmelCase     (							_UpperCamelCase		:	Dict						)		->				str:
 '''simple docstring'''


 return {"yes": True, "no": False}[value.lower()]




class   UpperCamelCase			(  argparse.RawDescriptionHelpFormatter ):

    def UpperCamelCase      (      self			:   Union[str, Any]				,     snake_case__			:   Tuple				,     snake_case__			:   Optional[Any]				,     snake_case__			:   List[str]				,     snake_case__			:   Dict     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      super()._format_usage(snake_case__				,     snake_case__				,     snake_case__				,     snake_case__     )
     SCREAMING_SNAKE_CASE		    =      usage.replace('<command> [<args>] '				,     ''     )
     return usage

