


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


def lowerCAmelCase_      (				UpperCamelCase__	: Optional[Any]				,   UpperCamelCase__	: Optional[int]=False				,   UpperCamelCase__	: Any=False				,   UpperCamelCase__	: Any=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 lowerCAmelCase_      (				UpperCamelCase__	: str				,   UpperCamelCase__	: List[str] ):






 """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 lowerCAmelCase_      (				UpperCamelCase__	: str ):






 """simple docstring"""







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


def lowerCAmelCase_      (				UpperCamelCase__	: Tuple				,   UpperCamelCase__	: Optional[Any]				,   UpperCamelCase__	: Dict ):






 """simple docstring"""







 __lowercase									=	dct.pop(UpperCamelCase__ )
 __lowercase									=	val


@torch.no_grad()
def lowerCAmelCase_      (				UpperCamelCase__	: Optional[Any]				,   UpperCamelCase__	: Any ):






 """simple docstring"""







 __lowercase									=	ViltConfig(image_size=384				,   patch_size=32				,   tie_word_embeddings=UpperCamelCase__ )
 __lowercase									=	False
 __lowercase									=	False
 __lowercase									=	False
 __lowercase									=	False
 if "vqa" in checkpoint_url:
  __lowercase									=	True
  __lowercase									=	3129
  __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							,		__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=384 )
 __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, 30522] )
  __lowercase									=	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"
  __lowercase									=	outputs.logits[0, 4, :].argmax(-1 ).item()
  assert tokenizer.decode([predicted_id] ) == "cats"
 elif vqa_model:
  __lowercase									=	torch.Size([1, 3129] )
  __lowercase									=	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"
  __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__":
       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)




"""simple docstring"""

import os

import numpy
import onnx


def lowerCAmelCase_      (				UpperCamelCase__	: Optional[int]				,   UpperCamelCase__	: List[str] ):






 """simple docstring"""







 __lowercase									=	a.name
 __lowercase									=	b.name

 __lowercase									=	""""""
 __lowercase									=	""""""

 __lowercase									=	a == b

 __lowercase									=	name_a
 __lowercase									=	name_b

 return res


def lowerCAmelCase_      (				UpperCamelCase__	: Union[str, Any]				,   UpperCamelCase__	: str				,   UpperCamelCase__	: Dict ):






 """simple docstring"""







 for i, input_name in enumerate(node_proto.input ):
  if input_name == name:
   node_proto.input.insert(UpperCamelCase__				,   UpperCamelCase__ )
   node_proto.input.pop(i + 1 )

 if node_proto.op_type == "If":
  _graph_replace_input_with(node_proto.attribute[0].g				,   UpperCamelCase__				,   UpperCamelCase__ )
  _graph_replace_input_with(node_proto.attribute[1].g				,   UpperCamelCase__				,   UpperCamelCase__ )
 if node_proto.op_type == "Loop":
  _graph_replace_input_with(node_proto.attribute[0].g				,   UpperCamelCase__				,   UpperCamelCase__ )


def lowerCAmelCase_      (				UpperCamelCase__	: Optional[Any]				,   UpperCamelCase__	: Optional[int]				,   UpperCamelCase__	: str ):






 """simple docstring"""







 for n in graph_proto.node:
  _node_replace_input_with(UpperCamelCase__				,   UpperCamelCase__				,   UpperCamelCase__ )


def lowerCAmelCase_      (				UpperCamelCase__	: Tuple				,   UpperCamelCase__	: str				,   UpperCamelCase__	: Any ):






 """simple docstring"""







 __lowercase									=	list(model.graph.initializer )
 __lowercase									=	list(model_without_ext.graph.initializer )
 for i, ref_i in ind_to_replace:
  assert inits_with_data[i].name == inits[i].name
  assert inits_with_data[ref_i].name == inits[ref_i].name
  assert i > ref_i

  __lowercase									=	inits[i].name
  __lowercase									=	inits[ref_i].name

  model_without_ext.graph.initializer.remove(inits[i] )

  # for n in model.graph.node:
  _graph_replace_input_with(model_without_ext.graph				,   UpperCamelCase__				,   UpperCamelCase__ )


def lowerCAmelCase_      (				UpperCamelCase__	: List[Any] ):






 """simple docstring"""







 __lowercase									=	os.path.dirname(UpperCamelCase__ )
 __lowercase									=	os.path.basename(UpperCamelCase__ )

 __lowercase									=	onnx.load(os.path.join(UpperCamelCase__				,   UpperCamelCase__ ) )

 __lowercase									=	list(model.graph.initializer )

 __lowercase									=	set()
 __lowercase									=	{}
 __lowercase									=	[]

 __lowercase									=	0

 for i in range(len(UpperCamelCase__ ) ):
  if i in dup_set:
   continue

  for j in range(i + 1				,   len(UpperCamelCase__ ) ):
   if j in dup_set:
    continue
   if _is_equal_tensor_proto(inits[i]				,   inits[j] ):
    dup_set.add(UpperCamelCase__ )
    dup_set.add(UpperCamelCase__ )

    __lowercase									=	inits[j].data_type
    __lowercase									=	numpy.prod(inits[j].dims )
    if dtype == 1:
     mem_size *= 4
    elif dtype == 6:
     mem_size *= 4
    elif dtype == 7 or dtype == 11:
     mem_size *= 8
    else:
     print("""unexpected data type: """				,   UpperCamelCase__ )
    total_reduced_size += mem_size

    __lowercase									=	inits[i].name
    __lowercase									=	inits[j].name

    if name_i in dup_map:
     dup_map[name_i].append(UpperCamelCase__ )
    else:
     __lowercase									=	[name_j]
    ind_to_replace.append((j, i) )

 print("""total reduced size: """				,   total_reduced_size / 1024 / 1024 / 1024				,   """GB""" )

 __lowercase									=	sorted(UpperCamelCase__ )
 _remove_dup_initializers_from_model(UpperCamelCase__				,   UpperCamelCase__				,   UpperCamelCase__ )

 __lowercase									=	"""optimized_""" + model_file_name
 __lowercase									=	os.path.join(UpperCamelCase__				,   UpperCamelCase__ )
 onnx.save(UpperCamelCase__				,   UpperCamelCase__ )

 return new_model
