


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()
snake_case    :   int	    =					logging.get_logger(__name__)
def     __lowerCamelCase						(			UpperCAmelCase_						:		Tuple			,  UpperCAmelCase_						:		Tuple=False			,  UpperCAmelCase_						:		List[str]=False			,  UpperCAmelCase_						:		int=False						):




      """simple docstring"""





      a						:Optional[int]										=	[]
      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						(			UpperCAmelCase_						:		Any			,  UpperCAmelCase_						:		int						):




      """simple docstring"""





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




      """simple docstring"""





      a						:Optional[Any]										=	['''head.weight''', '''head.bias''']
      for k in ignore_keys:
            state_dict.pop(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_						)
def     __lowerCamelCase						(			UpperCAmelCase_						:		Optional[Any]			,  UpperCAmelCase_						:		Optional[int]			,  UpperCAmelCase_						:		Dict						):




      """simple docstring"""





      a						:Optional[Any]										=	dct.pop(SCREAMING_SNAKE_CASE_						)
      a						:Optional[int]										=	val
@torch.no_grad()
def     __lowerCamelCase						(			UpperCAmelCase_						:		Any			,  UpperCAmelCase_						:		Optional[Any]						):




      """simple docstring"""





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

      # load state_dict of original model, remove and rename some keys
      a						:Union[str, Any]										=	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_			,  map_location='''cpu'''						)['''state_dict''']
      a						:int										=	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:
            a						:Tuple										=	['''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:
            a							,    a						:Optional[Any]										=	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
      a						:Any										=	ViltImageProcessor(size=384						)
      a						:Any										=	BertTokenizer.from_pretrained('''bert-base-uncased'''						)
      a						:int										=	ViltProcessor(SCREAMING_SNAKE_CASE_			,  SCREAMING_SNAKE_CASE_						)

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

      # Verify outputs
      if mlm_model:
            a						:List[str]										=	torch.Size([1, 11, 3_0522]						)
            a						:List[Any]										=	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"
            a						:Any										=	outputs.logits[0, 4, :].argmax(-1						).item()
            assert tokenizer.decode([predicted_id]						) == "cats"
      elif vqa_model:
            a						:Any										=	torch.Size([1, 3129]						)
            a						:Dict										=	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"
            a						:str										=	outputs.logits.argmax(-1						).item()
            assert model.config.idalabel[predicted_idx] == "2"
      elif nlvr_model:
            a						:List[Any]										=	torch.Size([1, 2]						)
            a						:int										=	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__":
     snake_case    :   Optional[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.'''
     )

     snake_case    :   Optional[Any]	    =					parser.parse_args()
     convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
     from PIL import Image

     from ..image_utils import load_image

if is_tf_available():
     import tensorflow as tf

     from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
     from ..tf_utils import stable_softmax

if is_torch_available():
     from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

snake_case    :   str	    =					logging.get_logger(__name__)







@add_end_docstrings(_snake_case							)
class 		_snake_case  (				_snake_case							):




 def __init__(		self ,  *_lowerCamelCase ,  **_lowerCamelCase   ):
       super().__init__(*_lowerCamelCase ,  **_lowerCamelCase   )
       requires_backends(self ,  '''vision'''   )
       self.check_model_type(
           TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
           if self.framework == '''tf'''
           else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING   )




 def 	SCREAMING_SNAKE_CASE__    (		self ,  _lowerCamelCase=None   ):
       a						:Tuple										=	{}
       if top_k is not None:
             a						:int										=	top_k
       return {}, {}, postprocess_params




 def __call__(		self ,  _lowerCamelCase ,  **_lowerCamelCase   ):
       return super().__call__(_lowerCamelCase ,  **_lowerCamelCase   )




 def 	SCREAMING_SNAKE_CASE__    (		self ,  _lowerCamelCase   ):
       a						:str										=	load_image(_lowerCamelCase   )
       a						:Any										=	self.image_processor(images=_lowerCamelCase ,  return_tensors=self.framework   )
       return model_inputs




 def 	SCREAMING_SNAKE_CASE__    (		self ,  _lowerCamelCase   ):
       a						:List[str]										=	self.model(**_lowerCamelCase   )
       return model_outputs




 def 	SCREAMING_SNAKE_CASE__    (		self ,  _lowerCamelCase ,  _lowerCamelCase=5   ):
       if top_k > self.model.config.num_labels:
             a						:List[Any]										=	self.model.config.num_labels

       if self.framework == "pt":
             a						:int										=	model_outputs.logits.softmax(-1   )[0]
             a							,    a						:Union[str, Any]										=	probs.topk(_lowerCamelCase   )
       elif self.framework == "tf":
             a						:Optional[Any]										=	stable_softmax(model_outputs.logits ,  axis=-1   )[0]
             a						:Union[str, Any]										=	tf.math.top_k(_lowerCamelCase ,  k=_lowerCamelCase   )
             a							,    a						:Optional[int]										=	topk.values.numpy(), topk.indices.numpy()
       else:
             raise ValueError(F'''Unsupported framework: {self.framework}'''   )

       a						:Optional[int]										=	scores.tolist()
       a						:str										=	ids.tolist()
       return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(_lowerCamelCase ,  _lowerCamelCase   )]





