

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







				"""simple docstring"""





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





								'''simple docstring'''







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

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


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





								'''simple docstring'''







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

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

								A__   				=       self.get_config()

								return config, pixel_values, labels


				def 						snake_case__					(				self			:	Tuple		)->   Optional[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=_a,initializer_range=self.initializer_range,encoder_stride=self.encoder_stride,)


				def 						snake_case__					(				self			:	Optional[Any],lowercase_			:	Tuple,lowercase_			:	Union[str, Any],lowercase_			:	int		)->   Union[str, Any]:





								'''simple docstring'''







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


				def 						snake_case__					(				self			:	Optional[Any],lowercase_			:	Optional[int],lowercase_			:	str,lowercase_			:	Tuple		)->   Optional[int]:





								'''simple docstring'''







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

								# test greyscale images
								A__   				=       1
								A__   				=       ViTForMaskedImageModeling(_a		)
								model.to(_a		)
								model.eval()

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


				def 						snake_case__					(				self			:	List[str],lowercase_			:	Optional[Any],lowercase_			:	Optional[int],lowercase_			:	List[str]		)->   str:





								'''simple docstring'''







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

								# test greyscale images
								A__   				=       1
								A__   				=       ViTForImageClassification(_a		)
								model.to(_a		)
								model.eval()

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


				def 						snake_case__					(				self			:	Dict		)->   str:





								'''simple docstring'''







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




@require_torch
class 			A						(							_UpperCAmelCase   ,							_UpperCAmelCase   ,							unittest.TestCase			):







				"""simple docstring"""





				lowerCamelCase		     =							(
				    (
				        ViTModel,
				        ViTForImageClassification,
				        ViTForMaskedImageModeling,
				    )
				    if is_torch_available()
				    else ()
				)
				lowerCamelCase		     =							(
				    {'feature-extraction': ViTModel, 'image-classification': ViTForImageClassification}
				    if is_torch_available()
				    else {}
				)
				lowerCamelCase		     =							True

				lowerCamelCase		     =							False
				lowerCamelCase		     =							False
				lowerCamelCase		     =							False


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





								'''simple docstring'''







								A__   				=       ViTModelTester(self		)
								A__   				=       ConfigTester(self,config_class=_a,has_text_modality=_a,hidden_size=3_7		)


				def 						snake_case__					(				self			:	int		)->   List[str]:





								'''simple docstring'''







								self.config_tester.run_common_tests()


				@unittest.skip(reason='ViT does not use inputs_embeds'		)
				def 						snake_case__					(				self			:	int		)->   Tuple:





								'''simple docstring'''







								pass


				def 						snake_case__					(				self			:	List[str]		)->   Tuple:





								'''simple docstring'''







								A__   				=       self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												A__   				=       model_class(_a		)
												self.assertIsInstance(model.get_input_embeddings(),(nn.Module)		)
												A__   				=       model.get_output_embeddings()
												self.assertTrue(x is None or isinstance(_a,nn.Linear		)		)


				def 						snake_case__					(				self			:	Any		)->   str:





								'''simple docstring'''







								A__   				=       self.model_tester.prepare_config_and_inputs_for_common()

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

												A__   				=       ["pixel_values"]
												self.assertListEqual(arg_names[:1],_a		)


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





								'''simple docstring'''







								A__   				=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_a		)


				def 						snake_case__					(				self			:	Dict		)->   Tuple:





								'''simple docstring'''







								A__   				=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_image_modeling(*_a		)


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





								'''simple docstring'''







								A__   				=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_a		)


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





								'''simple docstring'''







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




def 						_snake_case(   )							->				Tuple:





				'''simple docstring'''






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


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







				"""simple docstring"""





				@cached_property
				def 						snake_case__					(				self			:	int		)->   Optional[Any]:





								'''simple docstring'''







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


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





								'''simple docstring'''







								A__   				=       ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'		).to(_a		)

								A__   				=       self.default_image_processor
								A__   				=       prepare_img()
								A__   				=       image_processor(images=_a,return_tensors='pt'		).to(_a		)

								# forward pass
								with torch.no_grad():
												A__   				=       model(**_a		)

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

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

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


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





								'''simple docstring'''







								A__   				=       ViTModel.from_pretrained('facebook/dino-vits8'		).to(_a		)

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

								# forward pass
								with torch.no_grad():
												A__   				=       model(_a,interpolate_pos_encoding=_a		)

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

								A__   				=       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(_a		)

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


				@slow
				@require_accelerate
				@require_torch_gpu
				def 						snake_case__					(				self			:	int		)->   Union[str, Any]:





								'''simple docstring'''







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

								A__   				=       prepare_img()
								A__   				=       image_processor(images=_a,return_tensors='pt'		)
								A__   				=       inputs.pixel_values.to(_a		)

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




import numpy as np

from transformers import Pipeline
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     int				)							->				int:





    '''simple docstring'''






    A__   				=       np.max(SCREAMING_SNAKE_CASE__   ,      axis=-1   ,      keepdims=SCREAMING_SNAKE_CASE__				)
    A__   				=       np.exp(outputs - maxes				)
    return shifted_exp / shifted_exp.sum(axis=-1   ,      keepdims=SCREAMING_SNAKE_CASE__				)


class 			A						(							_UpperCAmelCase			):







    """simple docstring"""





    def 						snake_case__					(				self			:	Dict,**lowercase_			:	Tuple		)->   Tuple:





        '''simple docstring'''







        A__   				=       {}
        if "second_text" in kwargs:
            A__   				=       kwargs['second_text']
        return preprocess_kwargs, {}, {}


    def 						snake_case__					(				self			:	List[Any],lowercase_			:	int,lowercase_			:	Optional[int]=None		)->   List[str]:





        '''simple docstring'''







        return self.tokenizer(lowercase_,text_pair=lowercase_,return_tensors=self.framework		)


    def 						snake_case__					(				self			:	str,lowercase_			:	Dict		)->   List[str]:





        '''simple docstring'''







        return self.model(**lowercase_		)


    def 						snake_case__					(				self			:	Dict,lowercase_			:	Optional[int]		)->   Dict:





        '''simple docstring'''







        A__   				=       model_outputs.logits[0].numpy()
        A__   				=       softmax(lowercase_		)

        A__   				=       np.argmax(lowercase_		)
        A__   				=       self.model.config.idalabel[best_class]
        A__   				=       probabilities[best_class].item()
        A__   				=       logits.tolist()
        return {"label": label, "score": score, "logits": logits}

