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


				'''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 			A							(  self     :   Dict     ):


				'''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=lowercase      ,  initializer_range=self.initializer_range      ,  encoder_stride=self.encoder_stride      ,  )



	def 			A							(  self     :   Optional[Any]      ,  lowercase     :   Optional[int]      ,  lowercase     :   Tuple      ,  lowercase     :   List[Any]     ):


				'''simple docstring'''


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



	def 			A							(  self     :   str      ,  lowercase     :   Any      ,  lowercase     :   Optional[int]      ,  lowercase     :   Tuple     ):


				'''simple docstring'''


				_snake_case          =      ViTForMaskedImageModeling(config=lowercase     )
				model.to(lowercase     )
				model.eval()
				_snake_case          =      model(lowercase     )
				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(lowercase     )
				model.to(lowercase     )
				model.eval()

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



	def 			A							(  self     :   Dict      ,  lowercase     :   str      ,  lowercase     :   Optional[Any]      ,  lowercase     :   Union[str, Any]     ):


				'''simple docstring'''


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

				# test greyscale images
				_snake_case          =      1
				_snake_case          =      ViTForImageClassification(lowercase     )
				model.to(lowercase     )
				model.eval()

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





	def 			A							(  self     :   List[Any]     ):


				'''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__					(     UpperCAmelCase							,UpperCAmelCase							,unittest.TestCase ):





	'''simple docstring'''



	_UpperCAmelCase	:		Any          = (
	    (
	        ViTModel,
	        ViTForImageClassification,
	        ViTForMaskedImageModeling,
	    )
	    if is_torch_available()
	    else ()
	)
	_UpperCAmelCase	:		Tuple          = (
	    {"feature-extraction": ViTModel, "image-classification": ViTForImageClassification}
	    if is_torch_available()
	    else {}
	)
	_UpperCAmelCase	:		int          = True

	_UpperCAmelCase	:		Union[str, Any]          = False
	_UpperCAmelCase	:		List[str]          = False
	_UpperCAmelCase	:		Any          = False



	def 			A							(  self     :   Union[str, Any]     ):


				'''simple docstring'''


				_snake_case          =      ViTModelTester(self     )
				_snake_case          =      ConfigTester(self      ,  config_class=lowercase      ,  has_text_modality=lowercase      ,  hidden_size=37     )



	def 			A							(  self     :   List[str]     ):


				'''simple docstring'''


				self.config_tester.run_common_tests()



	@unittest.skip(reason='ViT does not use inputs_embeds'     )
	def 			A							(  self     :   Optional[Any]     ):


				'''simple docstring'''


				pass



	def 			A							(  self     :   Any     ):


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



	def 			A							(  self     :   Union[str, Any]     ):


				'''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(lowercase     )
							_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]      ,  lowercase     )



	def 			A							(  self     :   List[Any]     ):


				'''simple docstring'''


				_snake_case          =      self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowercase     )



	def 			A							(  self     :   List[str]     ):


				'''simple docstring'''


				_snake_case          =      self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_image_modeling(*lowercase     )



	def 			A							(  self     :   List[Any]     ):


				'''simple docstring'''


				_snake_case          =      self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*lowercase     )





	@slow
	def 			A							(  self     :   str     ):


				'''simple docstring'''


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






def 	a_		(	)			->  Optional[Any]:
			_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 			A							(  self     :   Tuple     ):


				'''simple docstring'''


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



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


				'''simple docstring'''


				_snake_case          =      ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'     ).to(lowercase     )

				_snake_case          =      self.default_image_processor
				_snake_case          =      prepare_img()
				_snake_case          =      image_processor(images=lowercase      ,  return_tensors='pt'     ).to(lowercase     )

				# forward pass
				with torch.no_grad():
							_snake_case          =      model(**lowercase     )

				# verify the logits
				_snake_case          =      torch.Size((1, 1_000)     )
				self.assertEqual(outputs.logits.shape      ,  lowercase     )

				_snake_case          =      torch.tensor([-0.2744, 0.8215, -0.0836]     ).to(lowercase     )

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



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


				'''simple docstring'''


				_snake_case          =      ViTModel.from_pretrained('facebook/dino-vits8'     ).to(lowercase     )

				_snake_case          =      ViTImageProcessor.from_pretrained('facebook/dino-vits8'      ,  size=480     )
				_snake_case          =      prepare_img()
				_snake_case          =      image_processor(images=lowercase      ,  return_tensors='pt'     )
				_snake_case          =      inputs.pixel_values.to(lowercase     )

				# forward pass
				with torch.no_grad():
							_snake_case          =      model(lowercase      ,  interpolate_pos_encoding=lowercase     )

				# verify the logits
				_snake_case          =      torch.Size((1, 3_601, 384)     )
				self.assertEqual(outputs.last_hidden_state.shape      ,  lowercase     )

				_snake_case          =      torch.tensor(
				    [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]     ).to(lowercase     )

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





	@slow
	@require_accelerate
	@require_torch_gpu
	def 			A							(  self     :   Optional[int]     ):


				'''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=lowercase      ,  return_tensors='pt'     )
				_snake_case          =      inputs.pixel_values.to(lowercase     )

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

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCamelCase	:    Tuple          =      logging.get_logger(__name__)

_lowerCamelCase	:    Union[str, Any]          =      {
    '''google/canine-s''': '''https://huggingface.co/google/canine-s/resolve/main/config.json''',
    # See all CANINE models at https://huggingface.co/models?filter=canine
}




class        SCREAMING_SNAKE_CASE__					(     UpperCAmelCase ):





	'''simple docstring'''



	_UpperCAmelCase	:		Any          = "canine"



	def __init__(  self     :   int      ,  lowercase     :   Optional[int]=768      ,  lowercase     :   Tuple=12      ,  lowercase     :   Union[str, Any]=12      ,  lowercase     :   Optional[int]=3_072      ,  lowercase     :   Tuple="gelu"      ,  lowercase     :   Optional[Any]=0.1      ,  lowercase     :   Tuple=0.1      ,  lowercase     :   int=16_384      ,  lowercase     :   Optional[int]=16      ,  lowercase     :   Optional[int]=0.02      ,  lowercase     :   Optional[Any]=1E-12      ,  lowercase     :   Optional[Any]=0      ,  lowercase     :   Dict=0xE000      ,  lowercase     :   Optional[Any]=0xE001      ,  lowercase     :   Union[str, Any]=4      ,  lowercase     :   str=4      ,  lowercase     :   Optional[int]=8      ,  lowercase     :   List[str]=16_384      ,  lowercase     :   Union[str, Any]=128      ,  **lowercase     :   Optional[Any]      ,  ):


				'''simple docstring'''


				super().__init__(pad_token_id=lowercase      ,  bos_token_id=lowercase      ,  eos_token_id=lowercase      ,  **lowercase     )

				_snake_case          =      max_position_embeddings
				_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          =      initializer_range
				_snake_case          =      type_vocab_size
				_snake_case          =      layer_norm_eps

				# Character config:
				_snake_case          =      downsampling_rate
				_snake_case          =      upsampling_kernel_size
				_snake_case          =      num_hash_functions
				_snake_case          =      num_hash_buckets
				_snake_case          =      local_transformer_stride