





import inspect
import unittest

from transformers import YolosConfig
from transformers.testing_utils import require_torch, 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import YolosForObjectDetection, YolosModel
		from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor






class 					lowercase__ :

			'''simple docstring'''


			def __init__(					self,	__magic_name__,	__magic_name__=13,	__magic_name__=[30, 30],	__magic_name__=2,	__magic_name__=3,	__magic_name__=True,	__magic_name__=True,	__magic_name__=32,	__magic_name__=5,	__magic_name__=4,	__magic_name__=37,	__magic_name__="gelu",	__magic_name__=0.1,	__magic_name__=0.1,	__magic_name__=10,	__magic_name__=0.02,	__magic_name__=3,	__magic_name__=None,	__magic_name__=8,	__magic_name__=10,	)				->   str:






							"""simple docstring"""
							UpperCamelCase__							:       Any      =   parent
							UpperCamelCase__							:       Union[str, Any]      =   batch_size
							UpperCamelCase__							:       Union[str, Any]      =   image_size
							UpperCamelCase__							:       Optional[Any]      =   patch_size
							UpperCamelCase__							:       Tuple      =   num_channels
							UpperCamelCase__							:       str      =   is_training
							UpperCamelCase__							:       List[str]      =   use_labels
							UpperCamelCase__							:       Tuple      =   hidden_size
							UpperCamelCase__							:       Any      =   num_hidden_layers
							UpperCamelCase__							:       int      =   num_attention_heads
							UpperCamelCase__							:       Tuple      =   intermediate_size
							UpperCamelCase__							:       Optional[Any]      =   hidden_act
							UpperCamelCase__							:       Dict      =   hidden_dropout_prob
							UpperCamelCase__							:       List[Any]      =   attention_probs_dropout_prob
							UpperCamelCase__							:       int      =   type_sequence_label_size
							UpperCamelCase__							:       Optional[int]      =   initializer_range
							UpperCamelCase__							:       str      =   num_labels
							UpperCamelCase__							:       Optional[Any]      =   scope
							UpperCamelCase__							:       Tuple      =   n_targets
							UpperCamelCase__							:       str      =   num_detection_tokens
							# we set the expected sequence length (which is used in several tests)
							# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
							UpperCamelCase__							:       List[str]      =   (image_size[1] // patch_size) * (image_size[0] // patch_size)
							UpperCamelCase__							:       Tuple      =   num_patches + 1 + self.num_detection_tokens




			def    UpperCamelCase__						(					self							)				->   Tuple:






							"""simple docstring"""
							UpperCamelCase__							:       Any      =   floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]							)

							UpperCamelCase__							:       Union[str, Any]      =   None
							if self.use_labels:
											# labels is a list of Dict (each Dict being the labels for a given example in the batch)
											UpperCamelCase__							:       Dict      =   []
											for i in range(self.batch_size							):
															UpperCamelCase__							:       str      =   {}
															UpperCamelCase__							:       Optional[Any]      =   torch.randint(
															    high=self.num_labels,	size=(self.n_targets,),	device=__magic_name__							)
															UpperCamelCase__							:       Optional[Any]      =   torch.rand(self.n_targets,	4,	device=__magic_name__							)
															labels.append(__magic_name__							)

							UpperCamelCase__							:       Optional[Any]      =   self.get_config()

							return config, pixel_values, labels




			def    UpperCamelCase__						(					self							)				->   Tuple:






							"""simple docstring"""
							return YolosConfig(
							    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=__magic_name__,	initializer_range=self.initializer_range,	num_detection_tokens=self.num_detection_tokens,	num_labels=self.num_labels,	)




			def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__,	__magic_name__							)				->   Any:






							"""simple docstring"""
							UpperCamelCase__							:       Dict      =   YolosModel(config=__magic_name__							)
							model.to(__magic_name__							)
							model.eval()
							UpperCamelCase__							:       Dict      =   model(__magic_name__							)
							self.parent.assertEqual(
							    result.last_hidden_state.shape,	(self.batch_size, self.expected_seq_len, self.hidden_size)							)




			def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__,	__magic_name__							)				->   Optional[Any]:






							"""simple docstring"""
							UpperCamelCase__							:       List[str]      =   YolosForObjectDetection(__magic_name__							)
							model.to(__magic_name__							)
							model.eval()

							UpperCamelCase__							:       Any      =   model(pixel_values=__magic_name__							)
							UpperCamelCase__							:       str      =   model(__magic_name__							)

							self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_detection_tokens, self.num_labels + 1)							)
							self.parent.assertEqual(result.pred_boxes.shape,	(self.batch_size, self.num_detection_tokens, 4)							)

							UpperCamelCase__							:       Optional[int]      =   model(pixel_values=__magic_name__,	labels=__magic_name__							)

							self.parent.assertEqual(result.loss.shape,	()							)
							self.parent.assertEqual(result.logits.shape,	(self.batch_size, self.num_detection_tokens, self.num_labels + 1)							)
							self.parent.assertEqual(result.pred_boxes.shape,	(self.batch_size, self.num_detection_tokens, 4)							)




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






							"""simple docstring"""
							UpperCamelCase__							:       List[str]      =   self.prepare_config_and_inputs()
							UpperCamelCase__							:       Union[str, Any]      =   config_and_inputs
							UpperCamelCase__							:       Optional[int]      =   {'''pixel_values''': pixel_values}
							return config, inputs_dict






@require_torch
class 					lowercase__ (			__lowerCamelCase    ,       __lowerCamelCase    ,       unittest.TestCase       ):

			'''simple docstring'''


			a		:							List[str]     	=	(YolosModel, YolosForObjectDetection) if is_torch_available() else ()
			a		:							Tuple     	=	(
			    {"feature-extraction": YolosModel, "object-detection": YolosForObjectDetection} if is_torch_available() else {}
			)

			a		:							Optional[int]     	=	False
			a		:							Tuple     	=	False
			a		:							Dict     	=	False
			a		:							Optional[int]     	=	False




			def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__,	__magic_name__=False							)				->   List[str]:






							"""simple docstring"""
							UpperCamelCase__							:       Optional[Any]      =   super()._prepare_for_class(__magic_name__,	__magic_name__,	return_labels=__magic_name__							)

							if return_labels:
											if model_class.__name__ == "YolosForObjectDetection":
															UpperCamelCase__							:       int      =   []
															for i in range(self.model_tester.batch_size							):
																			UpperCamelCase__							:       int      =   {}
																			UpperCamelCase__							:       int      =   torch.ones(
																			    size=(self.model_tester.n_targets,),	device=__magic_name__,	dtype=torch.long							)
																			UpperCamelCase__							:       Optional[int]      =   torch.ones(
																			    self.model_tester.n_targets,	4,	device=__magic_name__,	dtype=torch.float							)
																			labels.append(__magic_name__							)
															UpperCamelCase__							:       List[Any]      =   labels

							return inputs_dict




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






							"""simple docstring"""
							UpperCamelCase__							:       Tuple      =   YolosModelTester(self							)
							UpperCamelCase__							:       Optional[Any]      =   ConfigTester(self,	config_class=__magic_name__,	has_text_modality=__magic_name__,	hidden_size=37							)




			def    UpperCamelCase__						(					self							)				->   List[Any]:






							"""simple docstring"""
							self.config_tester.run_common_tests()




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






							"""simple docstring"""
							pass




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






							"""simple docstring"""
							UpperCamelCase__							:       Optional[Any]      =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
											UpperCamelCase__							:       List[Any]      =   model_class(__magic_name__							)
											self.assertIsInstance(model.get_input_embeddings(),	(nn.Module)							)
											UpperCamelCase__							:       Optional[int]      =   model.get_output_embeddings()
											self.assertTrue(x is None or isinstance(__magic_name__,	nn.Linear							)							)




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






							"""simple docstring"""
							UpperCamelCase__							:       Tuple      =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
											UpperCamelCase__							:       List[Any]      =   model_class(__magic_name__							)
											UpperCamelCase__							:       Tuple      =   inspect.signature(model.forward							)
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											UpperCamelCase__							:       Dict      =   [*signature.parameters.keys()]

											UpperCamelCase__							:       Union[str, Any]      =   ['''pixel_values''']
											self.assertListEqual(arg_names[:1],	__magic_name__							)




			def    UpperCamelCase__						(					self							)				->   Any:






							"""simple docstring"""
							UpperCamelCase__							:       Any      =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__magic_name__							)




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






							"""simple docstring"""
							UpperCamelCase__							:       Dict      =   self.model_tester.prepare_config_and_inputs_for_common()
							UpperCamelCase__							:       List[str]      =   True

							# in YOLOS, the seq_len is different
							UpperCamelCase__							:       List[str]      =   self.model_tester.expected_seq_len
							for model_class in self.all_model_classes:
											UpperCamelCase__							:       str      =   True
											UpperCamelCase__							:       Optional[int]      =   False
											UpperCamelCase__							:       Dict      =   True
											UpperCamelCase__							:       List[str]      =   model_class(__magic_name__							)
											model.to(__magic_name__							)
											model.eval()
											with torch.no_grad():
															UpperCamelCase__							:       str      =   model(**self._prepare_for_class(__magic_name__,	__magic_name__							)							)
											UpperCamelCase__							:       int      =   outputs.attentions
											self.assertEqual(len(__magic_name__							),	self.model_tester.num_hidden_layers							)

											# check that output_attentions also work using config
											del inputs_dict["output_attentions"]
											UpperCamelCase__							:       str      =   True
											UpperCamelCase__							:       Any      =   model_class(__magic_name__							)
											model.to(__magic_name__							)
											model.eval()
											with torch.no_grad():
															UpperCamelCase__							:       Dict      =   model(**self._prepare_for_class(__magic_name__,	__magic_name__							)							)
											UpperCamelCase__							:       List[str]      =   outputs.attentions
											self.assertEqual(len(__magic_name__							),	self.model_tester.num_hidden_layers							)

											self.assertListEqual(
											    list(attentions[0].shape[-3:]							),	[self.model_tester.num_attention_heads, seq_len, seq_len],	)
											UpperCamelCase__							:       List[Any]      =   len(__magic_name__							)

											# Check attention is always last and order is fine
											UpperCamelCase__							:       Optional[Any]      =   True
											UpperCamelCase__							:       Tuple      =   True
											UpperCamelCase__							:       str      =   model_class(__magic_name__							)
											model.to(__magic_name__							)
											model.eval()
											with torch.no_grad():
															UpperCamelCase__							:       int      =   model(**self._prepare_for_class(__magic_name__,	__magic_name__							)							)

											UpperCamelCase__							:       List[Any]      =   1
											self.assertEqual(out_len + added_hidden_states,	len(__magic_name__							)							)

											UpperCamelCase__							:       List[str]      =   outputs.attentions

											self.assertEqual(len(__magic_name__							),	self.model_tester.num_hidden_layers							)
											self.assertListEqual(
											    list(self_attentions[0].shape[-3:]							),	[self.model_tester.num_attention_heads, seq_len, seq_len],	)




			def    UpperCamelCase__						(					self							)				->   int:






							"""simple docstring"""
							def check_hidden_states_output(__magic_name__,	__magic_name__,	__magic_name__							):
											UpperCamelCase__							:       List[str]      =   model_class(__magic_name__							)
											model.to(__magic_name__							)
											model.eval()

											with torch.no_grad():
															UpperCamelCase__							:       Optional[int]      =   model(**self._prepare_for_class(__magic_name__,	__magic_name__							)							)

											UpperCamelCase__							:       List[Any]      =   outputs.hidden_states

											UpperCamelCase__							:       Any      =   getattr(
											    self.model_tester,	'''expected_num_hidden_layers''',	self.model_tester.num_hidden_layers + 1							)
											self.assertEqual(len(__magic_name__							),	__magic_name__							)

											# YOLOS has a different seq_length
											UpperCamelCase__							:       List[str]      =   self.model_tester.expected_seq_len

											self.assertListEqual(
											    list(hidden_states[0].shape[-2:]							),	[seq_length, self.model_tester.hidden_size],	)

							UpperCamelCase__							:       List[str]      =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
											UpperCamelCase__							:       Tuple      =   True
											check_hidden_states_output(__magic_name__,	__magic_name__,	__magic_name__							)

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											UpperCamelCase__							:       Optional[Any]      =   True

											check_hidden_states_output(__magic_name__,	__magic_name__,	__magic_name__							)




			def    UpperCamelCase__						(					self							)				->   List[str]:






							"""simple docstring"""
							UpperCamelCase__							:       Dict      =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_object_detection(*__magic_name__							)




			@slow
			def    UpperCamelCase__						(					self							)				->   Any:






							"""simple docstring"""
							for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											UpperCamelCase__							:       List[Any]      =   YolosModel.from_pretrained(__magic_name__							)
											self.assertIsNotNone(__magic_name__							)




def 	lowerCAmelCase_					(						)   ->			Any:
				UpperCamelCase__							:       int      =   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
				return image






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

			'''simple docstring'''


			@cached_property
			def    UpperCamelCase__						(					self							)				->   str:






							"""simple docstring"""
							return AutoImageProcessor.from_pretrained('''hustvl/yolos-small'''							) if is_vision_available() else None




			@slow
			def    UpperCamelCase__						(					self							)				->   str:






							"""simple docstring"""
							UpperCamelCase__							:       List[str]      =   YolosForObjectDetection.from_pretrained('''hustvl/yolos-small'''							).to(__magic_name__							)

							UpperCamelCase__							:       Optional[int]      =   self.default_image_processor
							UpperCamelCase__							:       Union[str, Any]      =   prepare_img()
							UpperCamelCase__							:       Optional[Any]      =   image_processor(images=__magic_name__,	return_tensors='''pt'''							).to(__magic_name__							)

							# forward pass
							with torch.no_grad():
											UpperCamelCase__							:       str      =   model(inputs.pixel_values							)

							# verify outputs
							UpperCamelCase__							:       Dict      =   torch.Size((1, 100, 92)							)
							self.assertEqual(outputs.logits.shape,	__magic_name__							)

							UpperCamelCase__							:       str      =   torch.tensor(
							    [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]],	device=__magic_name__,	)
							UpperCamelCase__							:       List[Any]      =   torch.tensor(
							    [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]],	device=__magic_name__							)
							self.assertTrue(torch.allclose(outputs.logits[0, :3, :3],	__magic_name__,	atol=1E-4							)							)
							self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3],	__magic_name__,	atol=1E-4							)							)

							# verify postprocessing
							UpperCamelCase__							:       List[str]      =   image_processor.post_process_object_detection(
							    __magic_name__,	threshold=0.3,	target_sizes=[image.size[::-1]]							)[0]
							UpperCamelCase__							:       Any      =   torch.tensor([0.9994, 0.9790, 0.9964, 0.9972, 0.9861]							).to(__magic_name__							)
							UpperCamelCase__							:       Dict      =   [75, 75, 17, 63, 17]
							UpperCamelCase__							:       int      =   torch.tensor([335.0609, 79.3848, 375.4216, 187.2495]							).to(__magic_name__							)

							self.assertEqual(len(results['''scores''']							),	5							)
							self.assertTrue(torch.allclose(results['''scores'''],	__magic_name__,	atol=1E-4							)							)
							self.assertSequenceEqual(results['''labels'''].tolist(),	__magic_name__							)
							self.assertTrue(torch.allclose(results['''boxes'''][0, :],	__magic_name__							)							)








import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_  							=				logging.get_logger(__name__)




def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[Any]	,     __UpperCAmelCase:    List[str]=False	,     __UpperCAmelCase:    List[Any]=False	,     __UpperCAmelCase:    int=False   )   ->			Optional[Any]:
    UpperCamelCase__							:       Union[str, Any]      =   []
    for i in range(config.num_hidden_layers   ):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append((f"transformer.blocks.{i}.norm1.weight", f"vilt.encoder.layer.{i}.layernorm_before.weight")   )
        rename_keys.append((f"transformer.blocks.{i}.norm1.bias", f"vilt.encoder.layer.{i}.layernorm_before.bias")   )
        rename_keys.append(
            (f"transformer.blocks.{i}.attn.proj.weight", f"vilt.encoder.layer.{i}.attention.output.dense.weight")   )
        rename_keys.append(
            (f"transformer.blocks.{i}.attn.proj.bias", f"vilt.encoder.layer.{i}.attention.output.dense.bias")   )
        rename_keys.append((f"transformer.blocks.{i}.norm2.weight", f"vilt.encoder.layer.{i}.layernorm_after.weight")   )
        rename_keys.append((f"transformer.blocks.{i}.norm2.bias", f"vilt.encoder.layer.{i}.layernorm_after.bias")   )
        rename_keys.append(
            (f"transformer.blocks.{i}.mlp.fc1.weight", f"vilt.encoder.layer.{i}.intermediate.dense.weight")   )
        rename_keys.append((f"transformer.blocks.{i}.mlp.fc1.bias", f"vilt.encoder.layer.{i}.intermediate.dense.bias")   )
        rename_keys.append((f"transformer.blocks.{i}.mlp.fc2.weight", f"vilt.encoder.layer.{i}.output.dense.weight")   )
        rename_keys.append((f"transformer.blocks.{i}.mlp.fc2.bias", f"vilt.encoder.layer.{i}.output.dense.bias")   )

    # embeddings
    rename_keys.extend(
        [
            # text embeddings
            ('''text_embeddings.word_embeddings.weight''', '''vilt.embeddings.text_embeddings.word_embeddings.weight'''),
            (
                '''text_embeddings.position_embeddings.weight''',
                '''vilt.embeddings.text_embeddings.position_embeddings.weight''',
            ),
            ('''text_embeddings.position_ids''', '''vilt.embeddings.text_embeddings.position_ids'''),
            (
                '''text_embeddings.token_type_embeddings.weight''',
                '''vilt.embeddings.text_embeddings.token_type_embeddings.weight''',
            ),
            ('''text_embeddings.LayerNorm.weight''', '''vilt.embeddings.text_embeddings.LayerNorm.weight'''),
            ('''text_embeddings.LayerNorm.bias''', '''vilt.embeddings.text_embeddings.LayerNorm.bias'''),
            # patch embeddings
            ('''transformer.cls_token''', '''vilt.embeddings.cls_token'''),
            ('''transformer.patch_embed.proj.weight''', '''vilt.embeddings.patch_embeddings.projection.weight'''),
            ('''transformer.patch_embed.proj.bias''', '''vilt.embeddings.patch_embeddings.projection.bias'''),
            ('''transformer.pos_embed''', '''vilt.embeddings.position_embeddings'''),
            # token type embeddings
            ('''token_type_embeddings.weight''', '''vilt.embeddings.token_type_embeddings.weight'''),
        ]   )

    # final layernorm + pooler
    rename_keys.extend(
        [
            ('''transformer.norm.weight''', '''vilt.layernorm.weight'''),
            ('''transformer.norm.bias''', '''vilt.layernorm.bias'''),
            ('''pooler.dense.weight''', '''vilt.pooler.dense.weight'''),
            ('''pooler.dense.bias''', '''vilt.pooler.dense.bias'''),
        ]   )

    # classifier head(s)
    if vqa_model:
        # classification head
        rename_keys.extend(
            [
                ('''vqa_classifier.0.weight''', '''classifier.0.weight'''),
                ('''vqa_classifier.0.bias''', '''classifier.0.bias'''),
                ('''vqa_classifier.1.weight''', '''classifier.1.weight'''),
                ('''vqa_classifier.1.bias''', '''classifier.1.bias'''),
                ('''vqa_classifier.3.weight''', '''classifier.3.weight'''),
                ('''vqa_classifier.3.bias''', '''classifier.3.bias'''),
            ]   )
    elif nlvr_model:
        # classification head
        rename_keys.extend(
            [
                ('''nlvr2_classifier.0.weight''', '''classifier.0.weight'''),
                ('''nlvr2_classifier.0.bias''', '''classifier.0.bias'''),
                ('''nlvr2_classifier.1.weight''', '''classifier.1.weight'''),
                ('''nlvr2_classifier.1.bias''', '''classifier.1.bias'''),
                ('''nlvr2_classifier.3.weight''', '''classifier.3.weight'''),
                ('''nlvr2_classifier.3.bias''', '''classifier.3.bias'''),
            ]   )
    else:
        pass

    return rename_keys




def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[Any]	,     __UpperCAmelCase:    List[Any]   )   ->			Optional[int]:
    for i in range(config.num_hidden_layers   ):
        UpperCamelCase__							:       Tuple      =   '''vilt.'''
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        UpperCamelCase__							:       Dict      =   state_dict.pop(f"transformer.blocks.{i}.attn.qkv.weight"   )
        UpperCamelCase__							:       Tuple      =   state_dict.pop(f"transformer.blocks.{i}.attn.qkv.bias"   )
        # next, add query, keys and values (in that order) to the state dict
        UpperCamelCase__							:       Union[str, Any]      =   in_proj_weight[
            : config.hidden_size, :
        ]
        UpperCamelCase__							:       Optional[Any]      =   in_proj_bias[: config.hidden_size]
        UpperCamelCase__							:       Tuple      =   in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        UpperCamelCase__							:       Any      =   in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        UpperCamelCase__							:       List[Any]      =   in_proj_weight[
            -config.hidden_size :, :
        ]
        UpperCamelCase__							:       List[Any]      =   in_proj_bias[-config.hidden_size :]




def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[Any]   )   ->			Any:
    UpperCamelCase__							:       int      =   ['''head.weight''', '''head.bias''']
    for k in ignore_keys:
        state_dict.pop(__UpperCAmelCase	,     __UpperCAmelCase   )




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Any	,     __UpperCAmelCase:    Dict	,     __UpperCAmelCase:    int   )   ->			List[str]:
    UpperCamelCase__							:       str      =   dct.pop(__UpperCAmelCase   )
    UpperCamelCase__							:       Optional[int]      =   val




@torch.no_grad()
def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[Any]	,     __UpperCAmelCase:    Union[str, Any]   )   ->			int:
    UpperCamelCase__							:       Any      =   ViltConfig(image_size=384	,     patch_size=32	,     tie_word_embeddings=__UpperCAmelCase   )
    UpperCamelCase__							:       Any      =   False
    UpperCamelCase__							:       Optional[int]      =   False
    UpperCamelCase__							:       List[Any]      =   False
    UpperCamelCase__							:       int      =   False
    if "vqa" in checkpoint_url:
        UpperCamelCase__							:       Any      =   True
        UpperCamelCase__							:       Optional[int]      =   3129
        UpperCamelCase__							:       Dict      =   '''huggingface/label-files'''
        UpperCamelCase__							:       Optional[Any]      =   '''vqa2-id2label.json'''
        UpperCamelCase__							:       Optional[int]      =   json.load(open(hf_hub_download(__UpperCAmelCase	,     __UpperCAmelCase	,     repo_type='''dataset'''   )	,     '''r'''   )   )
        UpperCamelCase__							:       List[str]      =   {int(__UpperCAmelCase   ): v for k, v in idalabel.items()}
        UpperCamelCase__							:       str      =   idalabel
        UpperCamelCase__							:       List[Any]      =   {v: k for k, v in idalabel.items()}
        UpperCamelCase__							:       Tuple      =   ViltForQuestionAnswering(__UpperCAmelCase   )
    elif "nlvr" in checkpoint_url:
        UpperCamelCase__							:       Optional[Any]      =   True
        UpperCamelCase__							:       Union[str, Any]      =   2
        UpperCamelCase__							:       int      =   {0: '''False''', 1: '''True'''}
        UpperCamelCase__							:       Optional[Any]      =   {v: k for k, v in config.idalabel.items()}
        UpperCamelCase__							:       Tuple      =   3
        UpperCamelCase__							:       Optional[Any]      =   ViltForImagesAndTextClassification(__UpperCAmelCase   )
    elif "irtr" in checkpoint_url:
        UpperCamelCase__							:       List[str]      =   True
        UpperCamelCase__							:       Any      =   ViltForImageAndTextRetrieval(__UpperCAmelCase   )
    elif "mlm_itm" in checkpoint_url:
        UpperCamelCase__							:       Any      =   True
        UpperCamelCase__							:       int      =   ViltForMaskedLM(__UpperCAmelCase   )
    else:
        raise ValueError('''Unknown model type'''   )

    # load state_dict of original model, remove and rename some keys
    UpperCamelCase__							:       int      =   torch.hub.load_state_dict_from_url(__UpperCAmelCase	,     map_location='''cpu'''   )['''state_dict''']
    UpperCamelCase__							:       Dict      =   create_rename_keys(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )
    for src, dest in rename_keys:
        rename_key(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )
    read_in_q_k_v(__UpperCAmelCase	,     __UpperCAmelCase   )
    if mlm_model or irtr_model:
        UpperCamelCase__							:       str      =   ['''itm_score.fc.weight''', '''itm_score.fc.bias''']
        for k in ignore_keys:
            state_dict.pop(__UpperCAmelCase	,     __UpperCAmelCase   )

    # load state dict into HuggingFace model
    model.eval()
    if mlm_model:
        UpperCamelCase__     ,UpperCamelCase__							:       Optional[Any]      =   model.load_state_dict(__UpperCAmelCase	,     strict=__UpperCAmelCase   )
        assert missing_keys == ["mlm_score.decoder.bias"]
    else:
        model.load_state_dict(__UpperCAmelCase   )

    # Define processor
    UpperCamelCase__							:       Union[str, Any]      =   ViltImageProcessor(size=384   )
    UpperCamelCase__							:       List[str]      =   BertTokenizer.from_pretrained('''bert-base-uncased'''   )
    UpperCamelCase__							:       Optional[int]      =   ViltProcessor(__UpperCAmelCase	,     __UpperCAmelCase   )

    # Forward pass on example inputs (image + text)
    if nlvr_model:
        UpperCamelCase__							:       Optional[int]      =   Image.open(requests.get('''https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'''	,     stream=__UpperCAmelCase   ).raw   )
        UpperCamelCase__							:       Union[str, Any]      =   Image.open(requests.get('''https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'''	,     stream=__UpperCAmelCase   ).raw   )
        UpperCamelCase__							:       Dict      =   (
            '''The left image contains twice the number of dogs as the right image, and at least two dogs in total are'''
            ''' standing.'''
        )
        UpperCamelCase__							:       str      =   processor(__UpperCAmelCase	,     __UpperCAmelCase	,     return_tensors='''pt'''   )
        UpperCamelCase__							:       Optional[Any]      =   processor(__UpperCAmelCase	,     __UpperCAmelCase	,     return_tensors='''pt'''   )
        UpperCamelCase__							:       Union[str, Any]      =   model(
            input_ids=encoding_a.input_ids	,     pixel_values=encoding_a.pixel_values	,     pixel_values_a=encoding_a.pixel_values	,     )
    else:
        UpperCamelCase__							:       Union[str, Any]      =   Image.open(requests.get('''http://images.cocodataset.org/val2017/000000039769.jpg'''	,     stream=__UpperCAmelCase   ).raw   )
        if mlm_model:
            UpperCamelCase__							:       int      =   '''a bunch of [MASK] laying on a [MASK].'''
        else:
            UpperCamelCase__							:       Optional[Any]      =   '''How many cats are there?'''
        UpperCamelCase__							:       List[Any]      =   processor(__UpperCAmelCase	,     __UpperCAmelCase	,     return_tensors='''pt'''   )
        UpperCamelCase__							:       Union[str, Any]      =   model(**__UpperCAmelCase   )

    # Verify outputs
    if mlm_model:
        UpperCamelCase__							:       str      =   torch.Size([1, 11, 3_0522]   )
        UpperCamelCase__							:       Optional[Any]      =   torch.tensor([-12.5061, -12.5123, -12.5174]   )
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )

        # verify masked token prediction equals "cats"
        UpperCamelCase__							:       Optional[Any]      =   outputs.logits[0, 4, :].argmax(-1   ).item()
        assert tokenizer.decode([predicted_id]   ) == "cats"
    elif vqa_model:
        UpperCamelCase__							:       List[Any]      =   torch.Size([1, 3129]   )
        UpperCamelCase__							:       Tuple      =   torch.tensor([-15.9495, -18.1472, -10.3041]   )
        assert torch.allclose(outputs.logits[0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )

        # verify vqa prediction equals "2"
        UpperCamelCase__							:       Any      =   outputs.logits.argmax(-1   ).item()
        assert model.config.idalabel[predicted_idx] == "2"
    elif nlvr_model:
        UpperCamelCase__							:       Dict      =   torch.Size([1, 2]   )
        UpperCamelCase__							:       str      =   torch.tensor([-2.8721, 2.1291]   )
        assert torch.allclose(outputs.logits[0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )
        assert outputs.logits.shape == expected_shape

    Path(__UpperCAmelCase   ).mkdir(exist_ok=__UpperCAmelCase   )
    print(f"Saving model and processor to {pytorch_dump_folder_path}"   )
    model.save_pretrained(__UpperCAmelCase   )
    processor.save_pretrained(__UpperCAmelCase   )


if __name__ == "__main__":
  UpperCAmelCase_  							=				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--checkpoint_url',
      default='https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt',
      type=str,
      help='URL of the checkpoint you\'d like to convert.',
  )
  parser.add_argument(
      '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
  )

  UpperCAmelCase_  							=				parser.parse_args()
  convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)

