




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			:





				'''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       , ):
											lowerCAmelCase   =		parent
											lowerCAmelCase   =		batch_size
											lowerCAmelCase   =		image_size
											lowerCAmelCase   =		patch_size
											lowerCAmelCase   =		num_channels
											lowerCAmelCase   =		is_training
											lowerCAmelCase   =		use_labels
											lowerCAmelCase   =		hidden_size
											lowerCAmelCase   =		num_hidden_layers
											lowerCAmelCase   =		num_attention_heads
											lowerCAmelCase   =		intermediate_size
											lowerCAmelCase   =		hidden_act
											lowerCAmelCase   =		hidden_dropout_prob
											lowerCAmelCase   =		attention_probs_dropout_prob
											lowerCAmelCase   =		type_sequence_label_size
											lowerCAmelCase   =		initializer_range
											lowerCAmelCase   =		scope
											lowerCAmelCase   =		encoder_stride

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



				def       __snake_case		( self						):
											lowerCAmelCase   =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

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

											lowerCAmelCase   =		self.get_config()

											return config, pixel_values, labels



				def       __snake_case		( self						):
											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       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		ViTModel(config=UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()
											lowerCAmelCase   =		model(UpperCAmelCase_						)
											self.parent.assertEqual(result.last_hidden_state.shape       , (self.batch_size, self.seq_length, self.hidden_size)						)



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		ViTForMaskedImageModeling(config=UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()
											lowerCAmelCase   =		model(UpperCAmelCase_						)
											self.parent.assertEqual(
											    result.reconstruction.shape       , (self.batch_size, self.num_channels, self.image_size, self.image_size)						)

											# test greyscale images
											lowerCAmelCase   =		1
											lowerCAmelCase   =		ViTForMaskedImageModeling(UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()

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



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.type_sequence_label_size
											lowerCAmelCase   =		ViTForImageClassification(UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()
											lowerCAmelCase   =		model(UpperCAmelCase_       , labels=UpperCAmelCase_						)
											self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.type_sequence_label_size)						)

											# test greyscale images
											lowerCAmelCase   =		1
											lowerCAmelCase   =		ViTForImageClassification(UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()

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






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




@require_torch
class 							__UpperCamelCase			( __UpperCAmelCase     ,					__UpperCAmelCase     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: List[str]										=(
				    (
				        ViTModel,
				        ViTForImageClassification,
				        ViTForMaskedImageModeling,
				    )
				    if is_torch_available()
				    else ()
				)
				__a							: int										=(
				    {"""feature-extraction""": ViTModel, """image-classification""": ViTForImageClassification}
				    if is_torch_available()
				    else {}
				)
				__a							: List[Any]										=True

				__a							: List[Any]										=False
				__a							: List[Any]										=False
				__a							: Any										=False



				def       __snake_case		( self						):
											lowerCAmelCase   =		ViTModelTester(self						)
											lowerCAmelCase   =		ConfigTester(self       , config_class=UpperCAmelCase_       , has_text_modality=UpperCAmelCase_       , hidden_size=37						)



				def       __snake_case		( self						):
											self.config_tester.run_common_tests()



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



				def       __snake_case		( self						):
											lowerCAmelCase   ,     lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs_for_common()

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



				def       __snake_case		( self						):
											lowerCAmelCase   ,     lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs_for_common()

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

																		lowerCAmelCase   =		['''pixel_values''']
																		self.assertListEqual(arg_names[:1]       , UpperCAmelCase_						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*UpperCAmelCase_						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase_						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_						)






				@slow
				def       __snake_case		( self						):
											for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowerCAmelCase   =		ViTModel.from_pretrained(UpperCAmelCase_						)
																		self.assertIsNotNone(UpperCAmelCase_						)




def       UpperCAmelCase       (     ):
							lowerCAmelCase   =		Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
							return image


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





				'''simple docstring'''



				@cached_property
				def       __snake_case		( self						):
											return ViTImageProcessor.from_pretrained('''google/vit-base-patch16-224'''						) if is_vision_available() else None



				@slow
				def       __snake_case		( self						):
											lowerCAmelCase   =		ViTForImageClassification.from_pretrained('''google/vit-base-patch16-224'''						).to(UpperCAmelCase_						)

											lowerCAmelCase   =		self.default_image_processor
											lowerCAmelCase   =		prepare_img()
											lowerCAmelCase   =		image_processor(images=UpperCAmelCase_       , return_tensors='''pt'''						).to(UpperCAmelCase_						)

											# forward pass
											with torch.no_grad():
																		lowerCAmelCase   =		model(**UpperCAmelCase_						)

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

											lowerCAmelCase   =		torch.tensor([-0.2744, 0.8215, -0.0836]						).to(UpperCAmelCase_						)

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



				@slow
				def       __snake_case		( self						):
											# ViT models have an `interpolate_pos_encoding` argument in their forward method,
											# allowing to interpolate the pre-trained position embeddings in order to use
											# the model on higher resolutions. The DINO model by Facebook AI leverages this
											# to visualize self-attention on higher resolution images.
											lowerCAmelCase   =		ViTModel.from_pretrained('''facebook/dino-vits8'''						).to(UpperCAmelCase_						)

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

											# forward pass
											with torch.no_grad():
																		lowerCAmelCase   =		model(UpperCAmelCase_       , interpolate_pos_encoding=UpperCAmelCase_						)

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

											lowerCAmelCase   =		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       __snake_case		( self						):
											lowerCAmelCase   =		ViTModel.from_pretrained('''facebook/dino-vits8'''       , torch_dtype=torch.floataa       , device_map='''auto'''						)
											lowerCAmelCase   =		self.default_image_processor

											lowerCAmelCase   =		prepare_img()
											lowerCAmelCase   =		image_processor(images=UpperCAmelCase_       , return_tensors='''pt'''						)
											lowerCAmelCase   =		inputs.pixel_values.to(UpperCAmelCase_						)

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









from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader


class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , UpperCAmelCase_ = True       , UpperCAmelCase_ = None       , UpperCAmelCase_ = False       , UpperCAmelCase_ = None       , UpperCAmelCase_ = True       , UpperCAmelCase_ = "arrow"       , **UpperCAmelCase_       , ):
											super().__init__(
											    split=UpperCAmelCase_       , features=UpperCAmelCase_       , cache_dir=UpperCAmelCase_       , keep_in_memory=UpperCAmelCase_       , streaming=UpperCAmelCase_       , **UpperCAmelCase_       , )
											lowerCAmelCase   =		load_from_cache_file
											lowerCAmelCase   =		file_format
											lowerCAmelCase   =		Spark(
											    df=UpperCAmelCase_       , features=UpperCAmelCase_       , cache_dir=UpperCAmelCase_       , working_dir=UpperCAmelCase_       , **UpperCAmelCase_       , )






				def       __snake_case		( self						):
											if self.streaming:
																		return self.builder.as_streaming_dataset(split=self.split						)
											lowerCAmelCase   =		None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
											self.builder.download_and_prepare(
											    download_mode=UpperCAmelCase_       , file_format=self._file_format       , )
											return self.builder.as_dataset(split=self.split						)



