
'''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__      :		str       							=  logging.get_logger(__name__)
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_=False    ,      SCREAMING_SNAKE_CASE_=False    ,      SCREAMING_SNAKE_CASE_=False						)     ->     int:




   """simple docstring"""
   _SCREAMING_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_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Union[str, Any]:




   """simple docstring"""
   for i in range(config.num_hidden_layers						):
      _SCREAMING_SNAKE_CASE      =				"""vilt."""
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      _SCREAMING_SNAKE_CASE      =				state_dict.pop(F"transformer.blocks.{i}.attn.qkv.weight"						)
      _SCREAMING_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
      _SCREAMING_SNAKE_CASE      =				in_proj_weight[
          : config.hidden_size, :
      ]
      _SCREAMING_SNAKE_CASE      =				in_proj_bias[: config.hidden_size]
      _SCREAMING_SNAKE_CASE      =				in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      _SCREAMING_SNAKE_CASE      =				in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      _SCREAMING_SNAKE_CASE      =				in_proj_weight[
          -config.hidden_size :, :
      ]
      _SCREAMING_SNAKE_CASE      =				in_proj_bias[-config.hidden_size :]
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_						)     ->     str:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Any:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				dct.pop(SCREAMING_SNAKE_CASE_						)
   _SCREAMING_SNAKE_CASE      =				val
@torch.no_grad()
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     str:




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

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

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

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

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

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

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

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

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


if __name__ == "__main__":
     UpperCamelCase__      :		List[Any]       							=  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__      :		str       							=  parser.parse_args()
     convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)








'''simple docstring'''



import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding

class    _a				(_lowerCamelCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE      =  ['image_processor', 'tokenizer']
       SCREAMING_SNAKE_CASE      =  'ChineseCLIPImageProcessor'
       SCREAMING_SNAKE_CASE      =  ('BertTokenizer', 'BertTokenizerFast')






       def __init__(							self       ,							A__=None       ,							A__=None       ,							**A__      )    ->			int:
          _SCREAMING_SNAKE_CASE      =				None
          if "feature_extractor" in kwargs:
             warnings.warn(
                 """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
                 """ instead."""       ,							A__       ,							)
             _SCREAMING_SNAKE_CASE      =				kwargs.pop("""feature_extractor"""      )

          _SCREAMING_SNAKE_CASE      =				image_processor if image_processor is not None else feature_extractor
          if image_processor is None:
             raise ValueError("""You need to specify an `image_processor`."""      )
          if tokenizer is None:
             raise ValueError("""You need to specify a `tokenizer`."""      )

          super().__init__(A__       ,							A__      )
          _SCREAMING_SNAKE_CASE      =				self.image_processor






       def __call__(							self       ,							A__=None       ,							A__=None       ,							A__=None       ,							**A__      )    ->			Optional[int]:

          if text is None and images is None:
             raise ValueError("""You have to specify either text or images. Both cannot be none."""      )

          if text is not None:
             _SCREAMING_SNAKE_CASE      =				self.tokenizer(A__       ,							return_tensors=A__       ,							**A__      )

          if images is not None:
             _SCREAMING_SNAKE_CASE      =				self.image_processor(A__       ,							return_tensors=A__       ,							**A__      )

          if text is not None and images is not None:
             _SCREAMING_SNAKE_CASE      =				image_features.pixel_values
             return encoding
          elif text is not None:
             return encoding
          else:
             return BatchEncoding(data=dict(**A__      )       ,							tensor_type=A__      )






       def        UpperCamelCase				(							self       ,							*A__       ,							**A__      )    ->			Dict:
          return self.tokenizer.batch_decode(*A__       ,							**A__      )






       def        UpperCamelCase				(							self       ,							*A__       ,							**A__      )    ->			Optional[Any]:
          return self.tokenizer.decode(*A__       ,							**A__      )






       @property
       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				self.tokenizer.model_input_names
          _SCREAMING_SNAKE_CASE      =				self.image_processor.model_input_names
          return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names      )      )






       @property
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          warnings.warn(
              """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,							A__       ,							)
          return self.image_processor_class






