

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






			"""simple docstring"""




			def __init__(     self			,					__UpperCamelCase			,					__UpperCamelCase=13			,					__UpperCamelCase=30			,					__UpperCamelCase=2			,					__UpperCamelCase=3			,					__UpperCamelCase=True			,					__UpperCamelCase=True			,					__UpperCamelCase=32			,					__UpperCamelCase=5			,					__UpperCamelCase=4			,					__UpperCamelCase=37			,					__UpperCamelCase="gelu"			,					__UpperCamelCase=0.1			,					__UpperCamelCase=0.1			,					__UpperCamelCase=10			,					__UpperCamelCase=0.02			,					__UpperCamelCase=None			,					__UpperCamelCase=2			,					):

				"""simple docstring"""


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

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


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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

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

				snake_case_   		=			self.get_config()

				return config, pixel_values, labels


			def       __lowerCAmelCase			(     self				):

				"""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			,					__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase				):

				"""simple docstring"""


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


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase				):

				"""simple docstring"""


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

				# test greyscale images
				snake_case_   		=			1
				snake_case_   		=			ViTForMaskedImageModeling(__UpperCamelCase				)
				model.to(__UpperCamelCase				)
				model.eval()

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


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase			,					__UpperCamelCase			,					__UpperCamelCase				):

				"""simple docstring"""


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

				# test greyscale images
				snake_case_   		=			1
				snake_case_   		=			ViTForImageClassification(__UpperCamelCase				)
				model.to(__UpperCamelCase				)
				model.eval()

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







			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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







@require_torch
class SCREAMING_SNAKE_CASE	(		__snake_case				,    __snake_case				,    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				):

				"""simple docstring"""


				snake_case_   		=			ViTModelTester(self				)
				snake_case_   		=			ConfigTester(self			,					config_class=__UpperCamelCase			,					has_text_modality=__UpperCamelCase			,					hidden_size=37				)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				self.config_tester.run_common_tests()


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

				"""simple docstring"""


				pass


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_		,							snake_case_   		=			self.model_tester.prepare_config_and_inputs_for_common()

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


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_		,							snake_case_   		=			self.model_tester.prepare_config_and_inputs_for_common()

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

					snake_case_   		=			['pixel_values']
					self.assertListEqual(arg_names[:1]			,					__UpperCamelCase				)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*__UpperCamelCase				)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_image_modeling(*__UpperCamelCase				)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase				)







			@slow
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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







def a():





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






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






			"""simple docstring"""




			@cached_property
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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


			@slow
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'				).to(__UpperCamelCase				)

				snake_case_   		=			self.default_image_processor
				snake_case_   		=			prepare_img()
				snake_case_   		=			image_processor(images=__UpperCamelCase			,					return_tensors='pt'				).to(__UpperCamelCase				)

				# forward pass
				with torch.no_grad():
					snake_case_   		=			model(**__UpperCamelCase				)

				# verify the logits
				snake_case_   		=			torch.Size((1, 10_00)				)
				self.assertEqual(outputs.logits.shape			,					__UpperCamelCase				)

				snake_case_   		=			torch.tensor([-0.2744, 0.8215, -0.0836]				).to(__UpperCamelCase				)

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


			@slow
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			ViTModel.from_pretrained('facebook/dino-vits8'				).to(__UpperCamelCase				)

				snake_case_   		=			ViTImageProcessor.from_pretrained('facebook/dino-vits8'			,					size=4_80				)
				snake_case_   		=			prepare_img()
				snake_case_   		=			image_processor(images=__UpperCamelCase			,					return_tensors='pt'				)
				snake_case_   		=			inputs.pixel_values.to(__UpperCamelCase				)

				# forward pass
				with torch.no_grad():
					snake_case_   		=			model(__UpperCamelCase			,					interpolate_pos_encoding=__UpperCamelCase				)

				# verify the logits
				snake_case_   		=			torch.Size((1, 36_01, 3_84)				)
				self.assertEqual(outputs.last_hidden_state.shape			,					__UpperCamelCase				)

				snake_case_   		=			torch.tensor(
				    [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]				).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				):

				"""simple docstring"""


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

				snake_case_   		=			prepare_img()
				snake_case_   		=			image_processor(images=__UpperCamelCase			,					return_tensors='pt'				)
				snake_case_   		=			inputs.pixel_values.to(__UpperCamelCase				)

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



# Copyright 2023 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.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


A									=  {'configuration_mra': ['MRA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MraConfig']}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A									=  [
		    'MRA_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'MraForMaskedLM',
		    'MraForMultipleChoice',
		    'MraForQuestionAnswering',
		    'MraForSequenceClassification',
		    'MraForTokenClassification',
		    'MraLayer',
		    'MraModel',
		    'MraPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_mra import MRA_PRETRAINED_CONFIG_ARCHIVE_MAP, MraConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_mra import (
				    MRA_PRETRAINED_MODEL_ARCHIVE_LIST,
				    MraForMaskedLM,
				    MraForMultipleChoice,
				    MraForQuestionAnswering,
				    MraForSequenceClassification,
				    MraForTokenClassification,
				    MraLayer,
				    MraModel,
				    MraPreTrainedModel,
				)
else:
		import sys

		A									=  _LazyModule(__name__, globals()['__file__'], _import_structure)
