


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				, _UpperCAmelCase   :						Optional[Any]				, _UpperCAmelCase   :						Optional[int]=13				, _UpperCAmelCase   :						Dict=30				, _UpperCAmelCase   :						List[Any]=2				, _UpperCAmelCase   :						str=3				, _UpperCAmelCase   :						Any=True				, _UpperCAmelCase   :						int=True				, _UpperCAmelCase   :						Dict=32				, _UpperCAmelCase   :						int=5				, _UpperCAmelCase   :						Optional[Any]=4				, _UpperCAmelCase   :						Dict=37				, _UpperCAmelCase   :						Optional[int]="gelu"				, _UpperCAmelCase   :						Optional[int]=0.1				, _UpperCAmelCase   :						Union[str, Any]=0.1				, _UpperCAmelCase   :						List[str]=10				, _UpperCAmelCase   :						int=0.02				, _UpperCAmelCase   :						Optional[int]=None				, _UpperCAmelCase   :						Optional[int]=2				, )						->       List[Any]:




							"""simple docstring"""


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

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

					def   lowerCamelCase__					(self   :						Tuple		)						->       int:




							"""simple docstring"""


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

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

							lowercase__					      =		self.get_config()

							return config, pixel_values, labels

					def   lowerCamelCase__					(self   :						List[str]		)						->       Union[str, 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=_UpperCAmelCase				, initializer_range=self.initializer_range				, encoder_stride=self.encoder_stride				, )

					def   lowerCamelCase__					(self   :						List[Any]				, _UpperCAmelCase   :						Dict				, _UpperCAmelCase   :						Any				, _UpperCAmelCase   :						Optional[Any]		)						->       Union[str, Any]:




							"""simple docstring"""


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

					def   lowerCamelCase__					(self   :						int				, _UpperCAmelCase   :						Optional[int]				, _UpperCAmelCase   :						str				, _UpperCAmelCase   :						Optional[Any]		)						->       Optional[int]:




							"""simple docstring"""


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

							# test greyscale images
							lowercase__					      =		1
							lowercase__					      =		ViTForMaskedImageModeling(_UpperCAmelCase		)
							model.to(_UpperCAmelCase		)
							model.eval()

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

					def   lowerCamelCase__					(self   :						Optional[Any]				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						Optional[Any]				, _UpperCAmelCase   :						str		)						->       List[str]:




							"""simple docstring"""


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

							# test greyscale images
							lowercase__					      =		1
							lowercase__					      =		ViTForImageClassification(_UpperCAmelCase		)
							model.to(_UpperCAmelCase		)
							model.eval()

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


					def   lowerCamelCase__					(self   :						int		)						->       Dict:




							"""simple docstring"""


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






@require_torch
class 							A   (	UpperCAmelCase__			,  UpperCAmelCase__			,  unittest.TestCase							):
					'''simple docstring'''






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

					A__      =			False
					A__      =			False
					A__      =			False

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




							"""simple docstring"""


							lowercase__					      =		ViTModelTester(self		)
							lowercase__					      =		ConfigTester(self				, config_class=_UpperCAmelCase				, has_text_modality=_UpperCAmelCase				, hidden_size=37		)

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




							"""simple docstring"""


							self.config_tester.run_common_tests()

					@unittest.skip(reason="""ViT does not use inputs_embeds"""		)
					def   lowerCamelCase__					(self   :						Optional[Any]		)						->       Dict:




							"""simple docstring"""


							pass

					def   lowerCamelCase__					(self   :						Dict		)						->       Any:




							"""simple docstring"""


							lowercase__      ,     lowercase__					      =		self.model_tester.prepare_config_and_inputs_for_common()

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

					def   lowerCamelCase__					(self   :						str		)						->       Tuple:




							"""simple docstring"""


							lowercase__      ,     lowercase__					      =		self.model_tester.prepare_config_and_inputs_for_common()

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

									lowercase__					      =		["""pixel_values"""]
									self.assertListEqual(arg_names[:1]				, _UpperCAmelCase		)

					def   lowerCamelCase__					(self   :						Tuple		)						->       List[Any]:




							"""simple docstring"""


							lowercase__					      =		self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_UpperCAmelCase		)

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




							"""simple docstring"""


							lowercase__					      =		self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_image_modeling(*_UpperCAmelCase		)

					def   lowerCamelCase__					(self   :						Optional[Any]		)						->       Optional[Any]:




							"""simple docstring"""


							lowercase__					      =		self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase		)


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




							"""simple docstring"""


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






def 			UpperCamelCase (     )       ->  List[str]:
		"""simple docstring"""




		lowercase__					      =		Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
		return image







@require_torch
@require_vision
class 							A   (	unittest.TestCase							):
					'''simple docstring'''






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




							"""simple docstring"""


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

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




							"""simple docstring"""


							lowercase__					      =		ViTForImageClassification.from_pretrained("""google/vit-base-patch16-224"""		).to(_UpperCAmelCase		)

							lowercase__					      =		self.default_image_processor
							lowercase__					      =		prepare_img()
							lowercase__					      =		image_processor(images=_UpperCAmelCase				, return_tensors="""pt"""		).to(_UpperCAmelCase		)

							# forward pass
							with torch.no_grad():
									lowercase__					      =		model(**_UpperCAmelCase		)

							# verify the logits
							lowercase__					      =		torch.Size((1, 1000)		)
							self.assertEqual(outputs.logits.shape				, _UpperCAmelCase		)

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

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

					@slow
					def   lowerCamelCase__					(self   :						Union[str, Any]		)						->       Union[str, Any]:




							"""simple docstring"""


							lowercase__					      =		ViTModel.from_pretrained("""facebook/dino-vits8"""		).to(_UpperCAmelCase		)

							lowercase__					      =		ViTImageProcessor.from_pretrained("""facebook/dino-vits8"""				, size=480		)
							lowercase__					      =		prepare_img()
							lowercase__					      =		image_processor(images=_UpperCAmelCase				, return_tensors="""pt"""		)
							lowercase__					      =		inputs.pixel_values.to(_UpperCAmelCase		)

							# forward pass
							with torch.no_grad():
									lowercase__					      =		model(_UpperCAmelCase				, interpolate_pos_encoding=_UpperCAmelCase		)

							# verify the logits
							lowercase__					      =		torch.Size((1, 3601, 384)		)
							self.assertEqual(outputs.last_hidden_state.shape				, _UpperCAmelCase		)

							lowercase__					      =		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(_UpperCAmelCase		)

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


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




							"""simple docstring"""


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

							lowercase__					      =		prepare_img()
							lowercase__					      =		image_processor(images=_UpperCAmelCase				, return_tensors="""pt"""		)
							lowercase__					      =		inputs.pixel_values.to(_UpperCAmelCase		)

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





class 							A   :
					'''simple docstring'''






					def __init__(self   :						int				, _UpperCAmelCase   :						list		)						->       None:




							"""simple docstring"""


							lowercase__					      =		set_counts
							lowercase__					      =		max(_UpperCAmelCase		)
							lowercase__					      =		len(_UpperCAmelCase		)
							lowercase__					      =		[1] * num_sets
							lowercase__					      =		list(range(_UpperCAmelCase		)		)

					def   lowerCamelCase__					(self   :						List[Any]				, _UpperCAmelCase   :						int				, _UpperCAmelCase   :						int		)						->       bool:




							"""simple docstring"""


							lowercase__					      =		self.get_parent(_UpperCAmelCase		)
							lowercase__					      =		self.get_parent(_UpperCAmelCase		)

							if src_parent == dst_parent:
									return False

							if self.ranks[dst_parent] >= self.ranks[src_parent]:
									self.set_counts[dst_parent] += self.set_counts[src_parent]
									lowercase__					      =		0
									lowercase__					      =		dst_parent
									if self.ranks[dst_parent] == self.ranks[src_parent]:
											self.ranks[dst_parent] += 1
									lowercase__					      =		self.set_counts[dst_parent]
							else:
									self.set_counts[src_parent] += self.set_counts[dst_parent]
									lowercase__					      =		0
									lowercase__					      =		src_parent
									lowercase__					      =		self.set_counts[src_parent]

							lowercase__					      =		max(self.max_set				, _UpperCAmelCase		)
							return True


					def   lowerCamelCase__					(self   :						List[Any]				, _UpperCAmelCase   :						int		)						->       int:




							"""simple docstring"""


							if self.parents[disj_set] == disj_set:
									return disj_set
							lowercase__					      =		self.get_parent(self.parents[disj_set]		)
							return self.parents[disj_set]

