"""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 _a			(		_snake_case				,	_snake_case=False				,	_snake_case=False				,	_snake_case=False							):







      """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			(		_snake_case				,	_snake_case							):







      """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			(		_snake_case							):







      """simple docstring"""






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

def _a			(		_snake_case				,	_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	dct.pop(_snake_case							)
      UpperCAmelCase          =	val

@torch.no_grad()
def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	ViltConfig(image_size=384				,	patch_size=32				,	tie_word_embeddings=_snake_case							)
      UpperCAmelCase          =	False
      UpperCAmelCase          =	False
      UpperCAmelCase          =	False
      UpperCAmelCase          =	False
      if "vqa" in checkpoint_url:
            UpperCAmelCase          =	True
            UpperCAmelCase          =	3129
            UpperCAmelCase          =	"""huggingface/label-files"""
            UpperCAmelCase          =	"""vqa2-id2label.json"""
            UpperCAmelCase          =	json.load(open(hf_hub_download(_snake_case				,	_snake_case				,	repo_type="""dataset"""							)				,	"""r"""							)							)
            UpperCAmelCase          =	{int(_snake_case							): v for k, v in idalabel.items()}
            UpperCAmelCase          =	idalabel
            UpperCAmelCase          =	{v: k for k, v in idalabel.items()}
            UpperCAmelCase          =	ViltForQuestionAnswering(_snake_case							)
      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(_snake_case							)
      elif "irtr" in checkpoint_url:
            UpperCAmelCase          =	True
            UpperCAmelCase          =	ViltForImageAndTextRetrieval(_snake_case							)
      elif "mlm_itm" in checkpoint_url:
            UpperCAmelCase          =	True
            UpperCAmelCase          =	ViltForMaskedLM(_snake_case							)
      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(_snake_case				,	map_location="""cpu"""							)["""state_dict"""]
      UpperCAmelCase          =	create_rename_keys(_snake_case				,	_snake_case				,	_snake_case				,	_snake_case							)
      for src, dest in rename_keys:
            rename_key(_snake_case				,	_snake_case				,	_snake_case							)
      read_in_q_k_v(_snake_case				,	_snake_case							)
      if mlm_model or irtr_model:
            UpperCAmelCase          =	["""itm_score.fc.weight""", """itm_score.fc.bias"""]
            for k in ignore_keys:
                  state_dict.pop(_snake_case				,	_snake_case							)

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

      # Define processor
      UpperCAmelCase          =	ViltImageProcessor(size=384							)
      UpperCAmelCase          =	BertTokenizer.from_pretrained("""bert-base-uncased"""							)
      UpperCAmelCase          =	ViltProcessor(_snake_case				,	_snake_case							)

      # 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=_snake_case							).raw							)
            UpperCAmelCase          =	Image.open(requests.get("""https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"""				,	stream=_snake_case							).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(_snake_case				,	_snake_case				,	return_tensors="""pt"""							)
            UpperCAmelCase          =	processor(_snake_case				,	_snake_case				,	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=_snake_case							).raw							)
            if mlm_model:
                  UpperCAmelCase          =	"""a bunch of [MASK] laying on a [MASK]."""
            else:
                  UpperCAmelCase          =	"""How many cats are there?"""
            UpperCAmelCase          =	processor(_snake_case				,	_snake_case				,	return_tensors="""pt"""							)
            UpperCAmelCase          =	model(**_snake_case							)

      # Verify outputs
      if mlm_model:
            UpperCAmelCase          =	torch.Size([1, 11, 3_0522]							)
            UpperCAmelCase          =	torch.tensor([-12.5061, -12.5123, -12.5174]							)
            assert outputs.logits.shape == expected_shape
            assert torch.allclose(outputs.logits[0, 0, :3]				,	_snake_case				,	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, 3129]							)
            UpperCAmelCase          =	torch.tensor([-15.9495, -18.1472, -10.3041]							)
            assert torch.allclose(outputs.logits[0, :3]				,	_snake_case				,	atol=1E-4							)
            assert outputs.logits.shape == expected_shape
            assert torch.allclose(outputs.logits[0, 0, :3]				,	_snake_case				,	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.8721, 2.1291]							)
            assert torch.allclose(outputs.logits[0, :3]				,	_snake_case				,	atol=1E-4							)
            assert outputs.logits.shape == expected_shape

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


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 numpy as np

def _a			(		_snake_case							):







      """simple docstring"""






      return 1 / (1 + np.exp(-vector							))


if __name__ == "__main__":
  import doctest

  doctest.testmod()



