





'''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()
_SCREAMING_SNAKE_CASE							     =			logging.get_logger(__name__)





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





def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[Any]		)		->  Union[str, Any]:
   snake_case 						=   ["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(__lowerCAmelCase      ,    __lowerCAmelCase		)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]      ,    __lowerCAmelCase       :							Optional[int]      ,    __lowerCAmelCase       :							str		)		->  str:
   snake_case 						=   dct.pop(__lowerCAmelCase		)
   snake_case 						=   val





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

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

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

   # Define processor
   snake_case 						=   ViltImageProcessor(size=3_84		)
   snake_case 						=   BertTokenizer.from_pretrained("""bert-base-uncased"""		)
   snake_case 						=   ViltProcessor(__lowerCAmelCase      ,    __lowerCAmelCase		)

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

   # Verify outputs
   if mlm_model:
      snake_case 						=   torch.Size([1, 11, 3_05_22]		)
      snake_case 						=   torch.tensor([-12.5061, -12.5123, -12.5174]		)
      assert outputs.logits.shape == expected_shape
      assert torch.allclose(outputs.logits[0, 0, :3]      ,    __lowerCAmelCase      ,    atol=1e-4		)

      # verify masked token prediction equals "cats"
      snake_case 						=   outputs.logits[0, 4, :].argmax(-1		).item()
      assert tokenizer.decode([predicted_id]		) == "cats"
   elif vqa_model:
      snake_case 						=   torch.Size([1, 31_29]		)
      snake_case 						=   torch.tensor([-15.9495, -18.1472, -10.3041]		)
      assert torch.allclose(outputs.logits[0, :3]      ,    __lowerCAmelCase      ,    atol=1e-4		)
      assert outputs.logits.shape == expected_shape
      assert torch.allclose(outputs.logits[0, 0, :3]      ,    __lowerCAmelCase      ,    atol=1e-4		)

      # verify vqa prediction equals "2"
      snake_case 						=   outputs.logits.argmax(-1		).item()
      assert model.config.idalabel[predicted_idx] == "2"
   elif nlvr_model:
      snake_case 						=   torch.Size([1, 2]		)
      snake_case 						=   torch.tensor([-2.8721, 2.1291]		)
      assert torch.allclose(outputs.logits[0, :3]      ,    __lowerCAmelCase      ,    atol=1e-4		)
      assert outputs.logits.shape == expected_shape

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


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

 _SCREAMING_SNAKE_CASE							     =			parser.parse_args()
 convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)








'''simple docstring'''




import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
 import optax
 from flax.training.common_utils import onehot

 from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
 from transformers.models.ta.modeling_flax_ta import shift_tokens_right
@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			_lowerCAmelCase		(			unittest.TestCase	):






     """simple docstring"""






     @slow
     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Optional[Any]:
        snake_case 						=   FlaxMTaForConditionalGeneration.from_pretrained("""google/mt5-small"""						)
        snake_case 						=   AutoTokenizer.from_pretrained("""google/mt5-small"""						)

        snake_case 						=   tokenizer("""Hello there"""   ,       return_tensors="""np"""						).input_ids
        snake_case 						=   tokenizer("""Hi I am"""   ,       return_tensors="""np"""						).input_ids

        snake_case 						=   shift_tokens_right(__snake_case   ,       model.config.pad_token_id   ,       model.config.decoder_start_token_id						)

        snake_case 						=   model(__snake_case   ,       decoder_input_ids=__snake_case						).logits
        snake_case 						=   optax.softmax_cross_entropy(__snake_case   ,       onehot(__snake_case   ,       logits.shape[-1]						)						).mean()

        snake_case 						=   -(labels.shape[-1] * loss.item())

        snake_case 						=   -84.91_27
        self.assertTrue(abs(mtf_score - EXPECTED_SCORE						) < 1e-4						)

