



'''simple docstring'''




import argparse
import os

from . import (
    ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    BART_PRETRAINED_MODEL_ARCHIVE_LIST,
    BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
    DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
    DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
    ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
    LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    T5_PRETRAINED_CONFIG_ARCHIVE_MAP,
    TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    XLM_PRETRAINED_CONFIG_ARCHIVE_MAP,
    XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
    AlbertConfig,
    BartConfig,
    BertConfig,
    CamembertConfig,
    CTRLConfig,
    DistilBertConfig,
    DPRConfig,
    ElectraConfig,
    FlaubertConfig,
    GPTaConfig,
    LayoutLMConfig,
    LxmertConfig,
    OpenAIGPTConfig,
    RobertaConfig,
    TaConfig,
    TFAlbertForPreTraining,
    TFBartForConditionalGeneration,
    TFBartForSequenceClassification,
    TFBertForPreTraining,
    TFBertForQuestionAnswering,
    TFBertForSequenceClassification,
    TFCamembertForMaskedLM,
    TFCTRLLMHeadModel,
    TFDistilBertForMaskedLM,
    TFDistilBertForQuestionAnswering,
    TFDPRContextEncoder,
    TFDPRQuestionEncoder,
    TFDPRReader,
    TFElectraForPreTraining,
    TFFlaubertWithLMHeadModel,
    TFGPTaLMHeadModel,
    TFLayoutLMForMaskedLM,
    TFLxmertForPreTraining,
    TFLxmertVisualFeatureEncoder,
    TFOpenAIGPTLMHeadModel,
    TFRobertaForCausalLM,
    TFRobertaForMaskedLM,
    TFRobertaForSequenceClassification,
    TFTaForConditionalGeneration,
    TFTransfoXLLMHeadModel,
    TFWavaVecaModel,
    TFXLMRobertaForMaskedLM,
    TFXLMWithLMHeadModel,
    TFXLNetLMHeadModel,
    TransfoXLConfig,
    WavaVecaConfig,
    WavaVecaModel,
    XLMConfig,
    XLMRobertaConfig,
    XLNetConfig,
    is_torch_available,
    load_pytorch_checkpoint_in_tfa_model,
)
from .utils import CONFIG_NAME, WEIGHTS_NAME, cached_file, logging


if is_torch_available():
						import numpy as np
						import torch

						from . import (
						    AlbertForPreTraining,
						    BartForConditionalGeneration,
						    BertForPreTraining,
						    BertForQuestionAnswering,
						    BertForSequenceClassification,
						    CamembertForMaskedLM,
						    CTRLLMHeadModel,
						    DistilBertForMaskedLM,
						    DistilBertForQuestionAnswering,
						    DPRContextEncoder,
						    DPRQuestionEncoder,
						    DPRReader,
						    ElectraForPreTraining,
						    FlaubertWithLMHeadModel,
						    GPTaLMHeadModel,
						    LayoutLMForMaskedLM,
						    LxmertForPreTraining,
						    LxmertVisualFeatureEncoder,
						    OpenAIGPTLMHeadModel,
						    RobertaForMaskedLM,
						    RobertaForSequenceClassification,
						    TaForConditionalGeneration,
						    TransfoXLLMHeadModel,
						    XLMRobertaForMaskedLM,
						    XLMWithLMHeadModel,
						    XLNetLMHeadModel,
						)


logging.set_verbosity_info()

lowerCamelCase							:Optional[int]       = {
    '''bart''': (
        BartConfig,
        TFBartForConditionalGeneration,
        TFBartForSequenceClassification,
        BartForConditionalGeneration,
        BART_PRETRAINED_MODEL_ARCHIVE_LIST,
    ),
    '''bert''': (
        BertConfig,
        TFBertForPreTraining,
        BertForPreTraining,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''bert-large-uncased-whole-word-masking-finetuned-squad''': (
        BertConfig,
        TFBertForQuestionAnswering,
        BertForQuestionAnswering,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''bert-large-cased-whole-word-masking-finetuned-squad''': (
        BertConfig,
        TFBertForQuestionAnswering,
        BertForQuestionAnswering,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''bert-base-cased-finetuned-mrpc''': (
        BertConfig,
        TFBertForSequenceClassification,
        BertForSequenceClassification,
        BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''dpr''': (
        DPRConfig,
        TFDPRQuestionEncoder,
        TFDPRContextEncoder,
        TFDPRReader,
        DPRQuestionEncoder,
        DPRContextEncoder,
        DPRReader,
        DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
    ),
    '''gpt2''': (
        GPTaConfig,
        TFGPTaLMHeadModel,
        GPTaLMHeadModel,
        GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''xlnet''': (
        XLNetConfig,
        TFXLNetLMHeadModel,
        XLNetLMHeadModel,
        XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''xlm''': (
        XLMConfig,
        TFXLMWithLMHeadModel,
        XLMWithLMHeadModel,
        XLM_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''xlm-roberta''': (
        XLMRobertaConfig,
        TFXLMRobertaForMaskedLM,
        XLMRobertaForMaskedLM,
        XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''transfo-xl''': (
        TransfoXLConfig,
        TFTransfoXLLMHeadModel,
        TransfoXLLMHeadModel,
        TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''openai-gpt''': (
        OpenAIGPTConfig,
        TFOpenAIGPTLMHeadModel,
        OpenAIGPTLMHeadModel,
        OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''roberta''': (
        RobertaConfig,
        TFRobertaForCausalLM,
        TFRobertaForMaskedLM,
        RobertaForMaskedLM,
        ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''layoutlm''': (
        LayoutLMConfig,
        TFLayoutLMForMaskedLM,
        LayoutLMForMaskedLM,
        LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
    ),
    '''roberta-large-mnli''': (
        RobertaConfig,
        TFRobertaForSequenceClassification,
        RobertaForSequenceClassification,
        ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''camembert''': (
        CamembertConfig,
        TFCamembertForMaskedLM,
        CamembertForMaskedLM,
        CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''flaubert''': (
        FlaubertConfig,
        TFFlaubertWithLMHeadModel,
        FlaubertWithLMHeadModel,
        FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''distilbert''': (
        DistilBertConfig,
        TFDistilBertForMaskedLM,
        DistilBertForMaskedLM,
        DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''distilbert-base-distilled-squad''': (
        DistilBertConfig,
        TFDistilBertForQuestionAnswering,
        DistilBertForQuestionAnswering,
        DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''lxmert''': (
        LxmertConfig,
        TFLxmertForPreTraining,
        LxmertForPreTraining,
        LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''lxmert-visual-feature-encoder''': (
        LxmertConfig,
        TFLxmertVisualFeatureEncoder,
        LxmertVisualFeatureEncoder,
        LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''ctrl''': (
        CTRLConfig,
        TFCTRLLMHeadModel,
        CTRLLMHeadModel,
        CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''albert''': (
        AlbertConfig,
        TFAlbertForPreTraining,
        AlbertForPreTraining,
        ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''t5''': (
        TaConfig,
        TFTaForConditionalGeneration,
        TaForConditionalGeneration,
        T5_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''electra''': (
        ElectraConfig,
        TFElectraForPreTraining,
        ElectraForPreTraining,
        ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
    '''wav2vec2''': (
        WavaVecaConfig,
        TFWavaVecaModel,
        WavaVecaModel,
        WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    ),
}



def      a      (			lowerCamelCase__						,				lowerCamelCase__						,				lowerCamelCase__						,				lowerCamelCase__						,				lowerCamelCase__=False						,				lowerCamelCase__=True   ):
			'''simple docstring'''





			if model_type not in MODEL_CLASSES:
						raise ValueError(f'Unrecognized model type, should be one of {list(MODEL_CLASSES.keys()   )}.'   )

			A_     :						Tuple     =					MODEL_CLASSES[model_type]

			# Initialise TF model
			if config_file in aws_config_map:
						A_     :						Tuple     =					cached_file(lowerCamelCase__						,				lowerCamelCase__						,				force_download=not use_cached_models   )
			A_     :						Dict     =					config_class.from_json_file(lowerCamelCase__   )
			A_     :						List[Any]     =					True
			A_     :						Tuple     =					True
			print(f'Building TensorFlow model from configuration: {config}'   )
			A_     :						Union[str, Any]     =					model_class(lowerCamelCase__   )

			# Load weights from tf checkpoint
			if pytorch_checkpoint_path in aws_config_map.keys():
						A_     :						List[str]     =					cached_file(
						    lowerCamelCase__						,				lowerCamelCase__						,				force_download=not use_cached_models   )
			# Load PyTorch checkpoint in tf2 model:
			A_     :						Optional[Any]     =					load_pytorch_checkpoint_in_tfa_model(lowerCamelCase__						,				lowerCamelCase__   )

			if compare_with_pt_model:
						A_     :						str     =					tf_model(tf_model.dummy_inputs						,				training=lowerCamelCase__   )  # build the network

						A_     :						List[Any]     =					torch.load(lowerCamelCase__						,				map_location="""cpu"""   )
						A_     :						List[Any]     =					pt_model_class.from_pretrained(
						    pretrained_model_name_or_path=lowerCamelCase__						,				config=lowerCamelCase__						,				state_dict=lowerCamelCase__   )

						with torch.no_grad():
									A_     :						List[Any]     =					pt_model(**pt_model.dummy_inputs   )

						A_     :						Optional[int]     =					pto[0].numpy()
						A_     :						str     =					tfo[0].numpy()
						A_     :						Union[str, Any]     =					np.amax(np.abs(np_pt - np_tf   )   )
						print(f'Max absolute difference between models outputs {diff}'   )
						assert diff <= 2E-2, f'Error, model absolute difference is >2e-2: {diff}'

			# Save pytorch-model
			print(f'Save TensorFlow model to {tf_dump_path}'   )
			tf_model.save_weights(lowerCamelCase__						,				save_format="""h5"""   )



def      a      (			lowerCamelCase__						,				lowerCamelCase__						,				lowerCamelCase__=None						,				lowerCamelCase__=None						,				lowerCamelCase__=False						,				lowerCamelCase__=False						,				lowerCamelCase__=False						,				lowerCamelCase__=False						,				):
			'''simple docstring'''





			if args_model_type is None:
						A_     :						List[str]     =					list(MODEL_CLASSES.keys()   )
			else:
						A_     :						Dict     =					[args_model_type]

			for j, model_type in enumerate(lowerCamelCase__						,				start=1   ):
						print("""=""" * 1_00   )
						print(f' Converting model type {j}/{len(lowerCamelCase__   )}: {model_type}'   )
						print("""=""" * 1_00   )
						if model_type not in MODEL_CLASSES:
									raise ValueError(f'Unrecognized model type {model_type}, should be one of {list(MODEL_CLASSES.keys()   )}.'   )

						A_     :						List[Any]     =					MODEL_CLASSES[model_type]

						if model_shortcut_names_or_path is None:
									A_     :						Tuple     =					list(aws_model_maps.keys()   )
						if config_shortcut_names_or_path is None:
									A_     :						int     =					model_shortcut_names_or_path

						for i, (model_shortcut_name, config_shortcut_name) in enumerate(
						    zip(lowerCamelCase__						,				lowerCamelCase__   )						,				start=1   ):
									print("""-""" * 1_00   )
									if "-squad" in model_shortcut_name or "-mrpc" in model_shortcut_name or "-mnli" in model_shortcut_name:
												if not only_convert_finetuned_models:
															print(f'    Skipping finetuned checkpoint {model_shortcut_name}'   )
															continue
												A_     :						Optional[int]     =					model_shortcut_name
									elif only_convert_finetuned_models:
												print(f'    Skipping not finetuned checkpoint {model_shortcut_name}'   )
												continue
									print(
									    f'    Converting checkpoint {i}/{len(lowerCamelCase__   )}: {model_shortcut_name} - model_type {model_type}'   )
									print("""-""" * 1_00   )

									if config_shortcut_name in aws_config_map:
												A_     :						int     =					cached_file(lowerCamelCase__						,				lowerCamelCase__						,				force_download=not use_cached_models   )
									else:
												A_     :						int     =					config_shortcut_name

									if model_shortcut_name in aws_model_maps:
												A_     :						List[Any]     =					cached_file(lowerCamelCase__						,				lowerCamelCase__						,				force_download=not use_cached_models   )
									else:
												A_     :						Optional[int]     =					model_shortcut_name

									if os.path.isfile(lowerCamelCase__   ):
												A_     :						Optional[int]     =					"""converted_model"""

									convert_pt_checkpoint_to_tf(
									    model_type=lowerCamelCase__						,				pytorch_checkpoint_path=lowerCamelCase__						,				config_file=lowerCamelCase__						,				tf_dump_path=os.path.join(lowerCamelCase__						,				model_shortcut_name + """-tf_model.h5"""   )						,				compare_with_pt_model=lowerCamelCase__						,				)
									if remove_cached_files:
												os.remove(lowerCamelCase__   )
												os.remove(lowerCamelCase__   )


if __name__ == "__main__":
						lowerCamelCase							:List[str]       = argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--tf_dump_path''', default=None, type=str, required=True, help='''Path to the output Tensorflow dump file.'''
						)
						parser.add_argument(
						    '''--model_type''',
						    default=None,
						    type=str,
						    help=(
						        F"Model type selected in the list of {list(MODEL_CLASSES.keys())}. If not given, will download and "
						        '''convert all the models from AWS.'''
						    ),
						)
						parser.add_argument(
						    '''--pytorch_checkpoint_path''',
						    default=None,
						    type=str,
						    help=(
						        '''Path to the PyTorch checkpoint path or shortcut name to download from AWS. '''
						        '''If not given, will download and convert all the checkpoints from AWS.'''
						    ),
						)
						parser.add_argument(
						    '''--config_file''',
						    default=None,
						    type=str,
						    help=(
						        '''The config json file corresponding to the pre-trained model. \n'''
						        '''This specifies the model architecture. If not given and '''
						        '''--pytorch_checkpoint_path is not given or is a shortcut name '''
						        '''use the configuration associated to the shortcut name on the AWS'''
						    ),
						)
						parser.add_argument(
						    '''--compare_with_pt_model''', action='''store_true''', help='''Compare Tensorflow and PyTorch model predictions.'''
						)
						parser.add_argument(
						    '''--use_cached_models''',
						    action='''store_true''',
						    help='''Use cached models if possible instead of updating to latest checkpoint versions.''',
						)
						parser.add_argument(
						    '''--remove_cached_files''',
						    action='''store_true''',
						    help='''Remove pytorch models after conversion (save memory when converting in batches).''',
						)
						parser.add_argument('''--only_convert_finetuned_models''', action='''store_true''', help='''Only convert finetuned models.''')
						lowerCamelCase							:Optional[Any]       = parser.parse_args()

						# if args.pytorch_checkpoint_path is not None:
						#     convert_pt_checkpoint_to_tf(args.model_type.lower(),
						#                                 args.pytorch_checkpoint_path,
						#                                 args.config_file if args.config_file is not None else args.pytorch_checkpoint_path,
						#                                 args.tf_dump_path,
						#                                 compare_with_pt_model=args.compare_with_pt_model,
						#                                 use_cached_models=args.use_cached_models)
						# else:
						convert_all_pt_checkpoints_to_tf(
						    args.model_type.lower() if args.model_type is not None else None,
						    args.tf_dump_path,
						    model_shortcut_names_or_path=[args.pytorch_checkpoint_path]
						    if args.pytorch_checkpoint_path is not None
						    else None,
						    config_shortcut_names_or_path=[args.config_file] if args.config_file is not None else None,
						    compare_with_pt_model=args.compare_with_pt_model,
						    use_cached_models=args.use_cached_models,
						    remove_cached_files=args.remove_cached_files,
						    only_convert_finetuned_models=args.only_convert_finetuned_models,
						)




'''simple docstring'''




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


	def __init__(self	,				lowercase	,				lowercase=13	,				lowercase=30	,				lowercase=2	,				lowercase=3	,				lowercase=True	,				lowercase=True	,				lowercase=32	,				lowercase=5	,				lowercase=4	,				lowercase=37	,				lowercase="gelu"	,				lowercase=0.1	,				lowercase=0.1	,				lowercase=10	,				lowercase=0.02	,				lowercase=None	,				lowercase=2	,				):
				A_     :						List[str]     =					parent
				A_     :						str     =					batch_size
				A_     :						Optional[Any]     =					image_size
				A_     :						List[str]     =					patch_size
				A_     :						List[str]     =					num_channels
				A_     :						List[str]     =					is_training
				A_     :						str     =					use_labels
				A_     :						List[str]     =					hidden_size
				A_     :						List[Any]     =					num_hidden_layers
				A_     :						Any     =					num_attention_heads
				A_     :						Any     =					intermediate_size
				A_     :						Optional[Any]     =					hidden_act
				A_     :						Optional[int]     =					hidden_dropout_prob
				A_     :						str     =					attention_probs_dropout_prob
				A_     :						Optional[int]     =					type_sequence_label_size
				A_     :						Any     =					initializer_range
				A_     :						int     =					scope
				A_     :						str     =					encoder_stride

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


	def    _a     (self							):
				A_     :						str     =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

				A_     :						Optional[Any]     =					None
				if self.use_labels:
							A_     :						str     =					ids_tensor([self.batch_size]	,				self.type_sequence_label_size							)

				A_     :						Union[str, Any]     =					self.get_config()

				return config, pixel_values, labels


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


	def    _a     (self	,				lowercase	,				lowercase	,				lowercase							):
				A_     :						List[str]     =					ViTModel(config=lowercase							)
				model.to(lowercase							)
				model.eval()
				A_     :						List[str]     =					model(lowercase							)
				self.parent.assertEqual(result.last_hidden_state.shape	,				(self.batch_size, self.seq_length, self.hidden_size)							)


	def    _a     (self	,				lowercase	,				lowercase	,				lowercase							):
				A_     :						List[str]     =					ViTForMaskedImageModeling(config=lowercase							)
				model.to(lowercase							)
				model.eval()
				A_     :						Tuple     =					model(lowercase							)
				self.parent.assertEqual(
				    result.reconstruction.shape	,				(self.batch_size, self.num_channels, self.image_size, self.image_size)							)

				# test greyscale images
				A_     :						Union[str, Any]     =					1
				A_     :						Any     =					ViTForMaskedImageModeling(lowercase							)
				model.to(lowercase							)
				model.eval()

				A_     :						str     =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
				A_     :						Optional[int]     =					model(lowercase							)
				self.parent.assertEqual(result.reconstruction.shape	,				(self.batch_size, 1, self.image_size, self.image_size)							)


	def    _a     (self	,				lowercase	,				lowercase	,				lowercase							):
				A_     :						Dict     =					self.type_sequence_label_size
				A_     :						str     =					ViTForImageClassification(lowercase							)
				model.to(lowercase							)
				model.eval()
				A_     :						List[str]     =					model(lowercase	,				labels=lowercase							)
				self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.type_sequence_label_size)							)

				# test greyscale images
				A_     :						Any     =					1
				A_     :						str     =					ViTForImageClassification(lowercase							)
				model.to(lowercase							)
				model.eval()

				A_     :						Optional[int]     =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
				A_     :						Union[str, Any]     =					model(lowercase							)
				self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.type_sequence_label_size)							)


	def    _a     (self							):
				A_     :						str     =					self.prepare_config_and_inputs()
				(
				    (
				    A_
				),   (
				    A_
				),   (
				    A_
				),   
				)     :						Optional[int]     =					config_and_inputs
				A_     :						Tuple     =					{"""pixel_values""": pixel_values}
				return config, inputs_dict







@require_torch
class      _lowerCAmelCase							(      __UpperCAmelCase				,      __UpperCAmelCase				,      unittest.TestCase			):
	__SCREAMING_SNAKE_CASE       :		Any		      = (
	    (
	        ViTModel,
	        ViTForImageClassification,
	        ViTForMaskedImageModeling,
	    )
	    if is_torch_available()
	    else ()
	)
	__SCREAMING_SNAKE_CASE       :		Union[str, Any]		      = (
	    {'feature-extraction': ViTModel, 'image-classification': ViTForImageClassification}
	    if is_torch_available()
	    else {}
	)
	__SCREAMING_SNAKE_CASE       :		Union[str, Any]		      = True

	__SCREAMING_SNAKE_CASE       :		Dict		      = False
	__SCREAMING_SNAKE_CASE       :		Optional[int]		      = False
	__SCREAMING_SNAKE_CASE       :		Optional[int]		      = False


	def    _a     (self							):
				A_     :						Any     =					ViTModelTester(self							)
				A_     :						str     =					ConfigTester(self	,				config_class=lowercase	,				has_text_modality=lowercase	,				hidden_size=37							)


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


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


	def    _a     (self							):
				A_,   A_     :						Union[str, Any]     =					self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
							A_     :						Union[str, Any]     =					model_class(lowercase							)
							self.assertIsInstance(model.get_input_embeddings()	,				(nn.Module)							)
							A_     :						Tuple     =					model.get_output_embeddings()
							self.assertTrue(x is None or isinstance(lowercase	,				nn.Linear							)							)


	def    _a     (self							):
				A_,   A_     :						List[Any]     =					self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
							A_     :						Optional[Any]     =					model_class(lowercase							)
							A_     :						Any     =					inspect.signature(model.forward							)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							A_     :						List[str]     =					[*signature.parameters.keys()]

							A_     :						Dict     =					["""pixel_values"""]
							self.assertListEqual(arg_names[:1]	,				lowercase							)


	def    _a     (self							):
				A_     :						Tuple     =					self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowercase							)


	def    _a     (self							):
				A_     :						List[Any]     =					self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_image_modeling(*lowercase							)


	def    _a     (self							):
				A_     :						Dict     =					self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*lowercase							)


	@slow
	def    _a     (self							):
				for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A_     :						Dict     =					ViTModel.from_pretrained(lowercase							)
							self.assertIsNotNone(lowercase							)






def      a      (			):
			'''simple docstring'''





			A_     :						List[Any]     =					Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
			return image







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


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


	@slow
	def    _a     (self							):
				A_     :						Optional[int]     =					ViTForImageClassification.from_pretrained("""google/vit-base-patch16-224"""							).to(lowercase							)

				A_     :						List[str]     =					self.default_image_processor
				A_     :						Tuple     =					prepare_img()
				A_     :						int     =					image_processor(images=lowercase	,				return_tensors="""pt"""							).to(lowercase							)

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

				# verify the logits
				A_     :						Dict     =					torch.Size((1, 1000)							)
				self.assertEqual(outputs.logits.shape	,				lowercase							)

				A_     :						str     =					torch.tensor([-0.27_44, 0.82_15, -0.08_36]							).to(lowercase							)

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


	@slow
	def    _a     (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.
				A_     :						Optional[int]     =					ViTModel.from_pretrained("""facebook/dino-vits8"""							).to(lowercase							)

				A_     :						List[Any]     =					ViTImageProcessor.from_pretrained("""facebook/dino-vits8"""	,				size=480							)
				A_     :						Dict     =					prepare_img()
				A_     :						str     =					image_processor(images=lowercase	,				return_tensors="""pt"""							)
				A_     :						int     =					inputs.pixel_values.to(lowercase							)

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

				# verify the logits
				A_     :						int     =					torch.Size((1, 3601, 384)							)
				self.assertEqual(outputs.last_hidden_state.shape	,				lowercase							)

				A_     :						List[Any]     =					torch.tensor(
				    [[4.23_40, 4.39_06, -6.66_92], [4.54_63, 1.89_28, -6.72_57], [4.44_29, 0.84_96, -5.85_85]]							).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							):
				A_     :						List[Any]     =					ViTModel.from_pretrained("""facebook/dino-vits8"""	,				torch_dtype=torch.floataa	,				device_map="""auto"""							)
				A_     :						int     =					self.default_image_processor

				A_     :						Any     =					prepare_img()
				A_     :						List[str]     =					image_processor(images=lowercase	,				return_tensors="""pt"""							)
				A_     :						Any     =					inputs.pixel_values.to(lowercase							)

				# forward pass to make sure inference works in fp16
				with torch.no_grad():
							A_     :						Optional[Any]     =					model(lowercase							)