


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()
a			:			List[str]												=					logging.get_logger(__name__)


def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase=False							,      _UpperCamelCase=False							,      _UpperCamelCase=False	):


			'''simple docstring'''







			__lowercase			       = []
			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 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			for i in range(config.num_hidden_layers	):
						__lowercase			       = '''vilt.'''
						# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
						__lowercase			       = state_dict.pop(F'transformer.blocks.{i}.attn.qkv.weight'	)
						__lowercase			       = state_dict.pop(F'transformer.blocks.{i}.attn.qkv.bias'	)
						# next, add query, keys and values (in that order) to the state dict
						__lowercase			       = in_proj_weight[
						    : config.hidden_size, :
						]
						__lowercase			       = in_proj_bias[: config.hidden_size]
						__lowercase			       = in_proj_weight[
						    config.hidden_size : config.hidden_size * 2, :
						]
						__lowercase			       = in_proj_bias[
						    config.hidden_size : config.hidden_size * 2
						]
						__lowercase			       = in_proj_weight[
						    -config.hidden_size :, :
						]
						__lowercase			       = in_proj_bias[-config.hidden_size :]


def 		lowercase_	(  _UpperCamelCase	):


			'''simple docstring'''







			__lowercase			       = ['''head.weight''', '''head.bias''']
			for k in ignore_keys:
						state_dict.pop(_UpperCamelCase							,      _UpperCamelCase	)


def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__lowercase			       = dct.pop(_UpperCamelCase	)
			__lowercase			       = val


@torch.no_grad()
def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__lowercase			       = ViltConfig(image_size=3_84							,      patch_size=32							,      tie_word_embeddings=_UpperCamelCase	)
			__lowercase			       = False
			__lowercase			       = False
			__lowercase			       = False
			__lowercase			       = False
			if "vqa" in checkpoint_url:
						__lowercase			       = True
						__lowercase			       = 31_29
						__lowercase			       = '''huggingface/label-files'''
						__lowercase			       = '''vqa2-id2label.json'''
						__lowercase			       = json.load(open(hf_hub_download(_UpperCamelCase							,      _UpperCamelCase							,      repo_type='''dataset'''	)							,      '''r'''	)	)
						__lowercase			       = {int(_UpperCamelCase	): v for k, v in idalabel.items()}
						__lowercase			       = idalabel
						__lowercase			       = {v: k for k, v in idalabel.items()}
						__lowercase			       = ViltForQuestionAnswering(_UpperCamelCase	)
			elif "nlvr" in checkpoint_url:
						__lowercase			       = True
						__lowercase			       = 2
						__lowercase			       = {0: '''False''', 1: '''True'''}
						__lowercase			       = {v: k for k, v in config.idalabel.items()}
						__lowercase			       = 3
						__lowercase			       = ViltForImagesAndTextClassification(_UpperCamelCase	)
			elif "irtr" in checkpoint_url:
						__lowercase			       = True
						__lowercase			       = ViltForImageAndTextRetrieval(_UpperCamelCase	)
			elif "mlm_itm" in checkpoint_url:
						__lowercase			       = True
						__lowercase			       = ViltForMaskedLM(_UpperCamelCase	)
			else:
						raise ValueError('''Unknown model type'''	)

			# load state_dict of original model, remove and rename some keys
			__lowercase			       = torch.hub.load_state_dict_from_url(_UpperCamelCase							,      map_location='''cpu'''	)['''state_dict''']
			__lowercase			       = 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:
						__lowercase			       = ['''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:
						__lowercase			       = model.load_state_dict(_UpperCamelCase							,      strict=_UpperCamelCase	)
						assert missing_keys == ["mlm_score.decoder.bias"]
			else:
						model.load_state_dict(_UpperCamelCase	)

			# Define processor
			__lowercase			       = ViltImageProcessor(size=3_84	)
			__lowercase			       = BertTokenizer.from_pretrained('''bert-base-uncased'''	)
			__lowercase			       = ViltProcessor(_UpperCamelCase							,      _UpperCamelCase	)

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

			# Verify outputs
			if mlm_model:
						__lowercase			       = torch.Size([1, 11, 3_05_22]	)
						__lowercase			       = torch.tensor([-12.50_61, -12.51_23, -12.51_74]	)
						assert outputs.logits.shape == expected_shape
						assert torch.allclose(outputs.logits[0, 0, :3]							,      _UpperCamelCase							,      atol=1E-4	)

						# verify masked token prediction equals "cats"
						__lowercase			       = outputs.logits[0, 4, :].argmax(-1	).item()
						assert tokenizer.decode([predicted_id]	) == "cats"
			elif vqa_model:
						__lowercase			       = torch.Size([1, 31_29]	)
						__lowercase			       = torch.tensor([-15.94_95, -18.14_72, -10.30_41]	)
						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"
						__lowercase			       = outputs.logits.argmax(-1	).item()
						assert model.config.idalabel[predicted_idx] == "2"
			elif nlvr_model:
						__lowercase			       = torch.Size([1, 2]	)
						__lowercase			       = 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__":
	a			:			Union[str, Any]												=					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.'''
	)

	a			:			Any												=					parser.parse_args()
	convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





'''simple docstring'''





import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCamelCase_    						=       logging.get_logger(__name__)






def 				_lowerCAmelCase			(   __magic_name__						:     List[str]						)				->  Dict:
						lowercase :							List[str]    		=R'''\w+[.]\d+'''
						lowercase :							List[str]    		=re.findall(__magic_name__   , __magic_name__						)
						for pat in pats:
												lowercase :							Optional[int]    		=key.replace(__magic_name__   , '''_'''.join(pat.split('''.'''						)						)						)
						return key






def 				_lowerCAmelCase			(   __magic_name__						:     List[str]   , __magic_name__						:     Union[str, Any]   , __magic_name__						:     Dict						)				->  str:
						lowercase :							Dict    		=pt_tuple_key[:-1] + ('''scale''',)
						if (
						    any('''norm''' in str_ for str_ in pt_tuple_key						)
						    and (pt_tuple_key[-1] == "bias")
						    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
						    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
						):
												lowercase :							str    		=pt_tuple_key[:-1] + ('''scale''',)
												return renamed_pt_tuple_key, pt_tensor
						elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
												lowercase :							str    		=pt_tuple_key[:-1] + ('''scale''',)
												return renamed_pt_tuple_key, pt_tensor

						# embedding
						if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
												lowercase :							Dict    		=pt_tuple_key[:-1] + ('''embedding''',)
												return renamed_pt_tuple_key, pt_tensor

						# conv layer
						lowercase :							Tuple    		=pt_tuple_key[:-1] + ('''kernel''',)
						if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
												lowercase :							Tuple    		=pt_tensor.transpose(2   , 3   , 1   , 0						)
												return renamed_pt_tuple_key, pt_tensor

						# linear layer
						lowercase :							str    		=pt_tuple_key[:-1] + ('''kernel''',)
						if pt_tuple_key[-1] == "weight":
												lowercase :							Optional[Any]    		=pt_tensor.T
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm weight
						lowercase :							Dict    		=pt_tuple_key[:-1] + ('''weight''',)
						if pt_tuple_key[-1] == "gamma":
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm bias
						lowercase :							Union[str, Any]    		=pt_tuple_key[:-1] + ('''bias''',)
						if pt_tuple_key[-1] == "beta":
												return renamed_pt_tuple_key, pt_tensor

						return pt_tuple_key, pt_tensor







def 				_lowerCAmelCase			(   __magic_name__						:     int   , __magic_name__						:     Optional[Any]   , __magic_name__						:     Union[str, Any]=42						)				->  List[str]:
						# Step 1: Convert pytorch tensor to numpy
						lowercase :							Optional[Any]    		={k: v.numpy() for k, v in pt_state_dict.items()}

						# Step 2: Since the model is stateless, get random Flax params
						lowercase :							str    		=flax_model.init_weights(PRNGKey(__magic_name__						)						)

						lowercase :							Dict    		=flatten_dict(__magic_name__						)
						lowercase :							Dict    		={}

						# Need to change some parameters name to match Flax names
						for pt_key, pt_tensor in pt_state_dict.items():
												lowercase :							Dict    		=rename_key(__magic_name__						)
												lowercase :							Optional[int]    		=tuple(renamed_pt_key.split('''.'''						)						)

												# Correctly rename weight parameters
												lowercase		,	lowercase :							Any    		=rename_key_and_reshape_tensor(__magic_name__   , __magic_name__   , __magic_name__						)

												if flax_key in random_flax_state_dict:
																		if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																								raise ValueError(
																								    f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
																								    f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''						)

        # also add unexpected weight so that warning is thrown
												lowercase :							Tuple    		=jnp.asarray(__magic_name__						)

						return unflatten_dict(__magic_name__						)





