

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()
_lowerCAmelCase  :      str			 =logging.get_logger(__name__)



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



def 			_A	(      SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	Dict									= ["head.weight", "head.bias"]
	for k in ignore_keys:
		state_dict.pop(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	UpperCAmelCase__:	int									= dct.pop(SCREAMING_SNAKE_CASE      )
	UpperCAmelCase__:	Any									= val



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

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

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

	# Define processor
	UpperCAmelCase__:	Optional[Any]									= ViltImageProcessor(size=3_8_4      )
	UpperCAmelCase__:	Tuple									= BertTokenizer.from_pretrained("bert-base-uncased"      )
	UpperCAmelCase__:	List[str]									= ViltProcessor(SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      )

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

	# Verify outputs
	if mlm_model:
		UpperCAmelCase__:	Optional[Any]									= torch.Size([1, 1_1, 3_0_5_2_2]      )
		UpperCAmelCase__:	Optional[int]									= 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]	,SCREAMING_SNAKE_CASE	,atol=1e-4      )

		# verify masked token prediction equals "cats"
		UpperCAmelCase__:	List[str]									= outputs.logits[0, 4, :].argmax(-1      ).item()
		assert tokenizer.decode([predicted_id]      ) == "cats"
	elif vqa_model:
		UpperCAmelCase__:	Any									= torch.Size([1, 3_1_2_9]      )
		UpperCAmelCase__:	str									= torch.tensor([-15.94_95, -18.14_72, -10.30_41]      )
		assert torch.allclose(outputs.logits[0, :3]	,SCREAMING_SNAKE_CASE	,atol=1e-4      )
		assert outputs.logits.shape == expected_shape
		assert torch.allclose(outputs.logits[0, 0, :3]	,SCREAMING_SNAKE_CASE	,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__:	Any									= torch.Size([1, 2]      )
		UpperCAmelCase__:	Tuple									= torch.tensor([-2.87_21, 2.12_91]      )
		assert torch.allclose(outputs.logits[0, :3]	,SCREAMING_SNAKE_CASE	,atol=1e-4      )
		assert outputs.logits.shape == expected_shape

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


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

					_lowerCAmelCase  :      List[Any]			 =parser.parse_args()
					convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


_lowerCAmelCase  :      int			 ="""
# Installazione di Transformers
! pip install transformers datasets
# Per installare dalla fonte invece dell'ultima versione rilasciata, commenta il comando sopra e
# rimuovi la modalità commento al comando seguente.
# ! pip install git+https://github.com/huggingface/transformers.git
"""

_lowerCAmelCase  :      List[str]			 =[{"""type""": """code""", """content""": INSTALL_CONTENT}]
_lowerCAmelCase  :      int			 ={
    """{processor_class}""": """FakeProcessorClass""",
    """{model_class}""": """FakeModelClass""",
    """{object_class}""": """FakeObjectClass""",
}