





'''simple docstring'''







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()
_lowercase       						=			logging.get_logger(__name__)

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

def    __UpperCamelCase	(	a    : int     ) ->List[str]:
	snake_case				=			['''head.weight''', '''head.bias''']
	for k in ignore_keys:
		state_dict.pop(a					,   a     )

def    __UpperCamelCase	(	a    : Dict					,   a    : List[str]					,   a    : List[Any]     ) ->List[str]:
	snake_case				=			dct.pop(a     )
	snake_case				=			val

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

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

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

	# Define processor
	snake_case				=			ViltImageProcessor(size=384     )
	snake_case				=			BertTokenizer.from_pretrained('''bert-base-uncased'''     )
	snake_case				=			ViltProcessor(a					,   a     )

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

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

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

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

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


if __name__ == "__main__":
			_lowercase       						=			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.'
			)

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








'''simple docstring'''







from __future__ import annotations

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import TFViTForImageClassification, TFViTModel


if is_vision_available():
   from PIL import Image

   from transformers import ViTImageProcessor
class 			_lowercase						:







     def __init__( self    ,   A__    ,   A__=13    ,   A__=30    ,   A__=2    ,   A__=3    ,   A__=True    ,   A__=True    ,   A__=32    ,   A__=2    ,   A__=4    ,   A__=37    ,   A__="gelu"    ,   A__=0.1    ,   A__=0.1    ,   A__=10    ,   A__=0.0_2    ,   A__=3    ,   A__=None    ,   ) ->	List[Any]:
      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

      # 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     UpperCamelCase			( self    ) ->	int:
      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     UpperCamelCase			( self    ) ->	int:
      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=A__    ,   initializer_range=self.initializer_range    ,   )







     def     UpperCamelCase			( self    ,   A__    ,   A__    ,   A__    ) ->	Union[str, Any]:
      snake_case				=			TFViTModel(config=A__    )
      snake_case				=			model(A__    ,   training=A__    )
      self.parent.assertEqual(result.last_hidden_state.shape    ,   (self.batch_size, self.seq_length, self.hidden_size)    )

      # Test with an image with different size than the one specified in config.
      snake_case				=			self.image_size // 2
      snake_case				=			pixel_values[:, :, :image_size, :image_size]
      snake_case				=			model(A__    ,   interpolate_pos_encoding=A__    ,   training=A__    )
      snake_case				=			(image_size // self.patch_size) ** 2 + 1
      self.parent.assertEqual(result.last_hidden_state.shape    ,   (self.batch_size, seq_length, self.hidden_size)    )







     def     UpperCamelCase			( self    ,   A__    ,   A__    ,   A__    ) ->	Optional[int]:
      snake_case				=			self.type_sequence_label_size
      snake_case				=			TFViTForImageClassification(A__    )
      snake_case				=			model(A__    ,   labels=A__    ,   training=A__    )
      self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.type_sequence_label_size)    )

      # Test with an image with different size than the one specified in config.
      snake_case				=			self.image_size // 2
      snake_case				=			pixel_values[:, :, :image_size, :image_size]
      snake_case				=			model(A__    ,   interpolate_pos_encoding=A__    ,   training=A__    )
      self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.type_sequence_label_size)    )

      # test greyscale images
      snake_case				=			1
      snake_case				=			TFViTForImageClassification(A__    )
      snake_case				=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
      snake_case				=			model(A__    )
      self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.type_sequence_label_size)    )







     def     UpperCamelCase			( self    ) ->	Union[str, Any]:
      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_tf
class 			_lowercase						(       __a				,		__a				,		unittest.TestCase					):
     _UpperCAmelCase										=					(TFViTModel, TFViTForImageClassification) if is_tf_available() else ()
     _UpperCAmelCase										=					(
         {'''feature-extraction''': TFViTModel, '''image-classification''': TFViTForImageClassification}
         if is_tf_available()
         else {}
     )

     _UpperCAmelCase										=					False
     _UpperCAmelCase										=					False
     _UpperCAmelCase										=					False







     def     UpperCamelCase			( self    ) ->	List[Any]:
      snake_case				=			TFViTModelTester(self    )
      snake_case				=			ConfigTester(self    ,   config_class=A__    ,   has_text_modality=A__    ,   hidden_size=37    )







     def     UpperCamelCase			( self    ) ->	int:
      self.config_tester.run_common_tests()







     @unittest.skip(reason='''ViT does not use inputs_embeds'''    )
     def     UpperCamelCase			( self    ) ->	int:
      pass







     @unittest.skip(reason='''ViT does not use inputs_embeds'''    )
     def     UpperCamelCase			( self    ) ->	str:
      pass







     def     UpperCamelCase			( self    ) ->	Union[str, Any]:
      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(A__    )
       self.assertIsInstance(model.get_input_embeddings()    ,   (tf.keras.layers.Layer)    )
       snake_case				=			model.get_output_embeddings()
       self.assertTrue(x is None or isinstance(A__    ,   tf.keras.layers.Layer    )    )







     def     UpperCamelCase			( self    ) ->	List[Any]:
      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(A__    )
       snake_case				=			inspect.signature(model.call    )
       # 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]    ,   A__    )







     def     UpperCamelCase			( self    ) ->	Union[str, Any]:
      snake_case				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*A__    )







     def     UpperCamelCase			( self    ) ->	Optional[Any]:
      snake_case				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*A__    )







     @slow
     def     UpperCamelCase			( self    ) ->	Any:
      snake_case				=			TFViTModel.from_pretrained('''google/vit-base-patch16-224'''    )
      self.assertIsNotNone(A__    )




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






@require_tf
@require_vision
class 			_lowercase						(       unittest.TestCase					):







     @cached_property
     def     UpperCamelCase			( self    ) ->	Optional[int]:
      return ViTImageProcessor.from_pretrained('''google/vit-base-patch16-224'''    ) if is_vision_available() else None







     @slow
     def     UpperCamelCase			( self    ) ->	Dict:
      snake_case				=			TFViTForImageClassification.from_pretrained('''google/vit-base-patch16-224'''    )

      snake_case				=			self.default_image_processor
      snake_case				=			prepare_img()
      snake_case				=			image_processor(images=A__    ,   return_tensors='''tf'''    )

      # forward pass
      snake_case				=			model(**A__    )

      # verify the logits
      snake_case				=			tf.TensorShape((1, 10_00)    )
      self.assertEqual(outputs.logits.shape    ,   A__    )

      snake_case				=			tf.constant([-0.2_7_4_4, 0.8_2_1_5, -0.0_8_3_6]    )

      tf.debugging.assert_near(outputs.logits[0, :3]    ,   A__    ,   atol=1e-4    )

