


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


def    _snake_case  (	A			)		->       List[str]:
							lowerCAmelCase__			  =      ['''head.weight''', '''head.bias''']
							for k in ignore_keys:
														state_dict.pop(A					,       A			)


def    _snake_case  (	A					,       A					,       A			)		->       str:
							lowerCAmelCase__			  =      dct.pop(A			)
							lowerCAmelCase__			  =      val


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

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

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

							# Define processor
							lowerCAmelCase__			  =      ViltImageProcessor(size=384			)
							lowerCAmelCase__			  =      BertTokenizer.from_pretrained('''bert-base-uncased'''			)
							lowerCAmelCase__			  =      ViltProcessor(A					,       A			)

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

							# Verify outputs
							if mlm_model:
														lowerCAmelCase__			  =      torch.Size([1, 11, 30522]			)
														lowerCAmelCase__			  =      torch.tensor([-12.5_061, -12.5_123, -12.5_174]			)
														assert outputs.logits.shape == expected_shape
														assert torch.allclose(outputs.logits[0, 0, :3]					,       A					,       atol=1E-4			)

														# verify masked token prediction equals "cats"
														lowerCAmelCase__			  =      outputs.logits[0, 4, :].argmax(-1			).item()
														assert tokenizer.decode([predicted_id]			) == "cats"
							elif vqa_model:
														lowerCAmelCase__			  =      torch.Size([1, 3129]			)
														lowerCAmelCase__			  =      torch.tensor([-15.9_495, -18.1_472, -10.3_041]			)
														assert torch.allclose(outputs.logits[0, :3]					,       A					,       atol=1E-4			)
														assert outputs.logits.shape == expected_shape
														assert torch.allclose(outputs.logits[0, 0, :3]					,       A					,       atol=1E-4			)

														# verify vqa prediction equals "2"
														lowerCAmelCase__			  =      outputs.logits.argmax(-1			).item()
														assert model.config.idalabel[predicted_idx] == "2"
							elif nlvr_model:
														lowerCAmelCase__			  =      torch.Size([1, 2]			)
														lowerCAmelCase__			  =      torch.tensor([-2.8_721, 2.1_291]			)
														assert torch.allclose(outputs.logits[0, :3]					,       A					,       atol=1E-4			)
														assert outputs.logits.shape == expected_shape

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


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)
import warnings

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




class 					_UpperCAmelCase			(     lowerCAmelCase  ):



 '''simple docstring'''




 __A			       =					['''image_processor''', '''tokenizer''']
 __A			       =					'''ViTImageProcessor'''
 __A			       =					('''CLIPTokenizer''', '''CLIPTokenizerFast''')







 def __init__(		self				:   Union[str, Any]							, lowercase_				:   Optional[int]=None							, lowercase_				:   Dict=None							, **lowercase_				:   List[str])						->	str:


   """simple docstring"""







   _UpperCamelCase           =    None
   if "feature_extractor" in kwargs:
     warnings.warn(
         "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
         " instead."							, lowercase_							, )
     _UpperCamelCase           =    kwargs.pop("feature_extractor")

   _UpperCamelCase           =    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__(lowercase_							, lowercase_)







 def __call__(		self				:   Any							, lowercase_				:   List[str]=None							, lowercase_				:   int=None							, lowercase_				:   Optional[Any]=None							, lowercase_				:   Dict=None							, **lowercase_				:   Dict)						->	int:


   """simple docstring"""







   if text is None and visual_prompt is None and images is None:
     raise ValueError("You have to specify either text, visual prompt or images.")

   if text is not None and visual_prompt is not None:
     raise ValueError("You have to specify exactly one type of prompt. Either text or visual prompt.")

   if text is not None:
     _UpperCamelCase           =    self.tokenizer(lowercase_							, return_tensors=lowercase_							, **lowercase_)

   if visual_prompt is not None:
     _UpperCamelCase           =    self.image_processor(lowercase_							, return_tensors=lowercase_							, **lowercase_)

   if images is not None:
     _UpperCamelCase           =    self.image_processor(lowercase_							, return_tensors=lowercase_							, **lowercase_)

   if visual_prompt is not None and images is not None:
     _UpperCamelCase           =    {
         "pixel_values": image_features.pixel_values,
         "conditional_pixel_values": prompt_features.pixel_values,
     }
     return encoding
   elif text is not None and images is not None:
     _UpperCamelCase           =    image_features.pixel_values
     return encoding
   elif text is not None:
     return encoding
   elif visual_prompt is not None:
     _UpperCamelCase           =    {
         "conditional_pixel_values": prompt_features.pixel_values,
     }
     return encoding
   else:
     return BatchEncoding(data=dict(**lowercase_)							, tensor_type=lowercase_)







 def 		__UpperCAmelCase				(		self				:   Optional[Any]							, *lowercase_				:   List[Any]							, **lowercase_				:   List[Any])						->	int:


   """simple docstring"""







   return self.tokenizer.batch_decode(*lowercase_							, **lowercase_)







 def 		__UpperCAmelCase				(		self				:   Union[str, Any]							, *lowercase_				:   Any							, **lowercase_				:   int)						->	Optional[int]:


   """simple docstring"""







   return self.tokenizer.decode(*lowercase_							, **lowercase_)







 @property
 def 		__UpperCAmelCase				(		self				:   Optional[int])						->	str:


   """simple docstring"""







   warnings.warn(
       "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."							, lowercase_							, )
   return self.image_processor_class







 @property
 def 		__UpperCAmelCase				(		self				:   int)						->	int:


   """simple docstring"""







   warnings.warn(
       "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."							, lowercase_							, )
   return self.image_processor




