
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__    :		Optional[Any]          =   logging.get_logger(__name__)

def 		UpperCamelCase	(      lowercase_ , lowercase_=False , lowercase_=False , lowercase_=False   )		->	Optional[int]:


  '''simple docstring'''






  lowercase__  :   str    		=      []
  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	(      lowercase_ , lowercase_   )		->	Optional[Any]:


  '''simple docstring'''






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

def 		UpperCamelCase	(      lowercase_   )		->	Tuple:


  '''simple docstring'''






  lowercase__  :   int    		=      ["""head.weight""", """head.bias"""]
  for k in ignore_keys:
    state_dict.pop(lowercase_ , lowercase_   )

def 		UpperCamelCase	(      lowercase_ , lowercase_ , lowercase_   )		->	List[Any]:


  '''simple docstring'''






  lowercase__  :   Optional[int]    		=      dct.pop(lowercase_   )
  lowercase__  :   Dict    		=      val

@torch.no_grad()
def 		UpperCamelCase	(      lowercase_ , lowercase_   )		->	Dict:


  '''simple docstring'''






  lowercase__  :   Optional[Any]    		=      ViltConfig(image_size=3_84 , patch_size=32 , tie_word_embeddings=lowercase_   )
  lowercase__  :   int    		=      False
  lowercase__  :   Optional[Any]    		=      False
  lowercase__  :   Optional[int]    		=      False
  lowercase__  :   Optional[Any]    		=      False
  if "vqa" in checkpoint_url:
    lowercase__  :   List[str]    		=      True
    lowercase__  :   Optional[Any]    		=      31_29
    lowercase__  :   str    		=      """huggingface/label-files"""
    lowercase__  :   Optional[Any]    		=      """vqa2-id2label.json"""
    lowercase__  :   Optional[Any]    		=      json.load(open(hf_hub_download(lowercase_ , lowercase_ , repo_type="""dataset"""   ) , """r"""   )   )
    lowercase__  :   int    		=      {int(lowercase_   ): v for k, v in idalabel.items()}
    lowercase__  :   List[str]    		=      idalabel
    lowercase__  :   Union[str, Any]    		=      {v: k for k, v in idalabel.items()}
    lowercase__  :   Union[str, Any]    		=      ViltForQuestionAnswering(lowercase_   )
  elif "nlvr" in checkpoint_url:
    lowercase__  :   Dict    		=      True
    lowercase__  :   Any    		=      2
    lowercase__  :   Optional[Any]    		=      {0: """False""", 1: """True"""}
    lowercase__  :   List[str]    		=      {v: k for k, v in config.idalabel.items()}
    lowercase__  :   List[str]    		=      3
    lowercase__  :   Any    		=      ViltForImagesAndTextClassification(lowercase_   )
  elif "irtr" in checkpoint_url:
    lowercase__  :   Optional[Any]    		=      True
    lowercase__  :   List[Any]    		=      ViltForImageAndTextRetrieval(lowercase_   )
  elif "mlm_itm" in checkpoint_url:
    lowercase__  :   str    		=      True
    lowercase__  :   str    		=      ViltForMaskedLM(lowercase_   )
  else:
    raise ValueError("""Unknown model type"""   )

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

  # Define processor
  lowercase__  :   Optional[int]    		=      ViltImageProcessor(size=3_84   )
  lowercase__  :   Union[str, Any]    		=      BertTokenizer.from_pretrained("""bert-base-uncased"""   )
  lowercase__  :   Optional[int]    		=      ViltProcessor(lowercase_ , lowercase_   )

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

  # Verify outputs
  if mlm_model:
    lowercase__  :   Union[str, Any]    		=      torch.Size([1, 11, 3_05_22]   )
    lowercase__  :   Tuple    		=      torch.tensor([-12.5061, -12.5123, -12.5174]   )
    assert outputs.logits.shape == expected_shape
    assert torch.allclose(outputs.logits[0, 0, :3] , lowercase_ , atol=1E-4   )

    # verify masked token prediction equals "cats"
    lowercase__  :   str    		=      outputs.logits[0, 4, :].argmax(-1   ).item()
    assert tokenizer.decode([predicted_id]   ) == "cats"
  elif vqa_model:
    lowercase__  :   Optional[int]    		=      torch.Size([1, 31_29]   )
    lowercase__  :   Dict    		=      torch.tensor([-15.9495, -18.1472, -10.3041]   )
    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"
    lowercase__  :   Any    		=      outputs.logits.argmax(-1   ).item()
    assert model.config.idalabel[predicted_idx] == "2"
  elif nlvr_model:
    lowercase__  :   List[str]    		=      torch.Size([1, 2]   )
    lowercase__  :   List[str]    		=      torch.tensor([-2.8721, 2.1291]   )
    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__":
     lowerCamelCase__    :		Optional[int]          =   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)








from __future__ import annotations

from collections.abc import Iterator
from typing import Any



class       _snake_case   :


 def __init__(					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__  :   Any    		=      data
   lowercase__  :   Node | None    		=      None



class       _snake_case   :


 def __init__(					self):


   '''simple docstring'''




   lowercase__  :   List[Any]    		=      None
   lowercase__  :   int    		=      None


 def __iter__(					self):


   '''simple docstring'''




   lowercase__  :   List[str]    		=      self.head
   while self.head:
     yield node.data
     lowercase__  :   str    		=      node.next
     if node == self.head:
       break


 def __len__(					self):


   '''simple docstring'''




   return sum(1 for _ in self)


 def __repr__(					self):


   '''simple docstring'''




   return "->".join(str(SCREAMING_SNAKE_CASE_) for item in iter(self))


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.insert_nth(len(self)					,  SCREAMING_SNAKE_CASE_)


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.insert_nth(0					,  SCREAMING_SNAKE_CASE_)


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   if index < 0 or index > len(self):
     raise IndexError("""list index out of range.""")
   lowercase__  :   Optional[int]    		=      Node(SCREAMING_SNAKE_CASE_)
   if self.head is None:
     lowercase__  :   Union[str, Any]    		=      new_node  # first node points itself
     lowercase__  :   Dict    		=      new_node
   elif index == 0:  # insert at head
     lowercase__  :   int    		=      self.head
     lowercase__  :   int    		=      new_node
   else:
     lowercase__  :   List[str]    		=      self.head
     for _ in range(index - 1):
       lowercase__  :   Union[str, Any]    		=      temp.next
     lowercase__  :   int    		=      temp.next
     lowercase__  :   Optional[Any]    		=      new_node
     if index == len(self) - 1:  # insert at tail
       lowercase__  :   Union[str, Any]    		=      new_node


 def      lowercase__     (					self):


   '''simple docstring'''




   return self.delete_nth(0)


 def      lowercase__     (					self):


   '''simple docstring'''




   return self.delete_nth(len(self) - 1)


 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_ = 0):


   '''simple docstring'''




   if not 0 <= index < len(self):
     raise IndexError("""list index out of range.""")
   lowercase__  :   str    		=      self.head
   if self.head == self.tail:  # just one node
     lowercase__  :   List[Any]    		=      None
   elif index == 0:  # delete head node
     lowercase__  :   str    		=      self.tail.next.next
     lowercase__  :   List[Any]    		=      self.head.next
   else:
     lowercase__  :   Optional[int]    		=      self.head
     for _ in range(index - 1):
       lowercase__  :   List[Any]    		=      temp.next
     lowercase__  :   Optional[Any]    		=      temp.next
     lowercase__  :   Dict    		=      temp.next.next
     if index == len(self) - 1:  # delete at tail
       lowercase__  :   Tuple    		=      temp
   return delete_node.data



 def      lowercase__     (					self):


   '''simple docstring'''




   return len(self) == 0

def 		UpperCamelCase	(      )		->	None:


  '''simple docstring'''






  lowercase__  :   int    		=      CircularLinkedList()
  assert len(lowercase_   ) == 0
  assert circular_linked_list.is_empty() is True
  assert str(lowercase_   ) == ""

  try:
    circular_linked_list.delete_front()
    raise AssertionError  # This should not happen
  except IndexError:
    assert True  # This should happen

  try:
    circular_linked_list.delete_tail()
    raise AssertionError  # This should not happen
  except IndexError:
    assert True  # This should happen

  try:
    circular_linked_list.delete_nth(-1   )
    raise AssertionError
  except IndexError:
    assert True

  try:
    circular_linked_list.delete_nth(0   )
    raise AssertionError
  except IndexError:
    assert True

  assert circular_linked_list.is_empty() is True
  for i in range(5   ):
    assert len(lowercase_   ) == i
    circular_linked_list.insert_nth(lowercase_ , i + 1   )
  assert str(lowercase_   ) == "->".join(str(lowercase_   ) for i in range(1 , 6   )   )

  circular_linked_list.insert_tail(6   )
  assert str(lowercase_   ) == "->".join(str(lowercase_   ) for i in range(1 , 7   )   )
  circular_linked_list.insert_head(0   )
  assert str(lowercase_   ) == "->".join(str(lowercase_   ) for i in range(0 , 7   )   )

  assert circular_linked_list.delete_front() == 0
  assert circular_linked_list.delete_tail() == 6
  assert str(lowercase_   ) == "->".join(str(lowercase_   ) for i in range(1 , 6   )   )
  assert circular_linked_list.delete_nth(2   ) == 3

  circular_linked_list.insert_nth(2 , 3   )
  assert str(lowercase_   ) == "->".join(str(lowercase_   ) for i in range(1 , 6   )   )

  assert circular_linked_list.is_empty() is False


if __name__ == "__main__":
     import doctest

     doctest.testmod()






