
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      :				Union[str, Any]          = logging.get_logger(__name__)


def    A (      lowercase					, lowercase=False					, lowercase=False					, lowercase=False				)     ->						Union[str, Any]:
		'''simple docstring'''

		UpperCamelCase						=  []
		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 (      lowercase					, lowercase				)     ->						Optional[Any]:
		'''simple docstring'''

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


def    A (      lowercase				)     ->						Union[str, Any]:
		'''simple docstring'''

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


def    A (      lowercase					, lowercase					, lowercase				)     ->						Any:
		'''simple docstring'''

		UpperCamelCase						=  dct.pop(lowercase				)
		UpperCamelCase						=  val


@torch.no_grad()
def    A (      lowercase					, lowercase				)     ->						Tuple:
		'''simple docstring'''

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

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

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

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

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

		# Verify outputs
		if mlm_model:
				UpperCamelCase						=  torch.Size([1, 11, 30_522]				)
				UpperCamelCase						=  torch.tensor([-1_2.5_0_6_1, -1_2.5_1_2_3, -1_2.5_1_7_4]				)
				assert outputs.logits.shape == expected_shape
				assert torch.allclose(outputs.logits[0, 0, :3]					, lowercase					, atol=1e-4				)

				# verify masked token prediction equals "cats"
				UpperCamelCase						=  outputs.logits[0, 4, :].argmax(-1				).item()
				assert tokenizer.decode([predicted_id]				) == "cats"
		elif vqa_model:
				UpperCamelCase						=  torch.Size([1, 3_129]				)
				UpperCamelCase						=  torch.tensor([-1_5.9_4_9_5, -1_8.1_4_7_2, -1_0.3_0_4_1]				)
				assert torch.allclose(outputs.logits[0, :3]					, lowercase					, atol=1e-4				)
				assert outputs.logits.shape == expected_shape
				assert torch.allclose(outputs.logits[0, 0, :3]					, lowercase					, atol=1e-4				)

				# verify vqa prediction equals "2"
				UpperCamelCase						=  outputs.logits.argmax(-1				).item()
				assert model.config.idalabel[predicted_idx] == "2"
		elif nlvr_model:
				UpperCamelCase						=  torch.Size([1, 2]				)
				UpperCamelCase						=  torch.tensor([-2.8_7_2_1, 2.1_2_9_1]				)
				assert torch.allclose(outputs.logits[0, :3]					, lowercase					, atol=1e-4				)
				assert outputs.logits.shape == expected_shape

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


if __name__ == "__main__":
					_UpperCAmelCase      :				List[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."
					)

					_UpperCAmelCase      :				Optional[int]          = parser.parse_args()
					convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


def    A (      lowercase				)     ->						str:
		'''simple docstring'''

		UpperCamelCase						=  int(lowercase				)
		if decimal in (0, 1):  # Exit cases for the recursion
				return str(lowercase				)
		UpperCamelCase			,						UpperCamelCase						=  divmod(lowercase					, 2				)
		return binary_recursive(lowercase				) + str(lowercase				)


def    A (      lowercase				)     ->						str:
		'''simple docstring'''

		UpperCamelCase						=  str(lowercase				).strip()
		if not number:
				raise ValueError('No input value was provided'				)
		UpperCamelCase						=  '-' if number.startswith('-'				) else ''
		UpperCamelCase						=  number.lstrip('-'				)
		if not number.isnumeric():
				raise ValueError('Input value is not an integer'				)
		return f'''{negative}0b{binary_recursive(int(lowercase				)				)}'''


if __name__ == "__main__":
					from doctest import testmod

					testmod()
