




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



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	int,					SCREAMING_SNAKE_CASE__:	Optional[int]=False,					SCREAMING_SNAKE_CASE__:	Any=False,					SCREAMING_SNAKE_CASE__:	Any=False    )  ->			Dict:

   """simple docstring"""
   __a				     =    []
   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	(				SCREAMING_SNAKE_CASE__:	Tuple,					SCREAMING_SNAKE_CASE__:	Any    )  ->			Tuple:

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



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	List[str]    )  ->			List[Any]:

   """simple docstring"""
   __a				     =    ["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(UpperCAmelCase__,					UpperCAmelCase__    )



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Tuple,					SCREAMING_SNAKE_CASE__:	List[str],					SCREAMING_SNAKE_CASE__:	Union[str, Any]    )  ->			Union[str, Any]:

   """simple docstring"""
   __a				     =    dct.pop(UpperCAmelCase__    )
   __a				     =    val



@torch.no_grad()
def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	str,					SCREAMING_SNAKE_CASE__:	Optional[int]    )  ->			List[str]:

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

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

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

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

   # Verify outputs
   if mlm_model:
      __a				     =    torch.Size([1, 11, 30522]    )
      __a				     =    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],					UpperCAmelCase__,					atol=1e-4    )

      # verify masked token prediction equals "cats"
      __a				     =    outputs.logits[0, 4, :].argmax(-1    ).item()
      assert tokenizer.decode([predicted_id]    ) == "cats"
   elif vqa_model:
      __a				     =    torch.Size([1, 3129]    )
      __a				     =    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],					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"
      __a				     =    outputs.logits.argmax(-1    ).item()
      assert model.config.idalabel[predicted_idx] == "2"
   elif nlvr_model:
      __a				     =    torch.Size([1, 2]    )
      __a				     =    torch.tensor([-2.8_7_2_1, 2.1_2_9_1]    )
      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      :       Dict      			=      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      :       int      			=      parser.parse_args()
 convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)







'''simple docstring'''



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	int    )  ->			bool:

   """simple docstring"""
   return number & 1 == 0


if __name__ == "__main__":
 import doctest

 doctest.testmod()