






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



def 					_lowerCamelCase(			lowerCAmelCase__				:		Tuple    ,    lowerCAmelCase__				:		Dict=False    ,    lowerCAmelCase__				:		List[str]=False    ,    lowerCAmelCase__				:		Union[str, Any]=False      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	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 					_lowerCamelCase(			lowerCAmelCase__				:		Any    ,    lowerCAmelCase__				:		Optional[Any]      ):





				'''simple docstring'''







				for i in range(config.num_hidden_layers      ):
								SCREAMING_SNAKE_CASE_   :	str													=   'vilt.'
								# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   state_dict.pop(F'''transformer.blocks.{i}.attn.qkv.weight'''      )
								SCREAMING_SNAKE_CASE_   :	str													=   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_   :	List[str]													=   in_proj_weight[
								    : config.hidden_size, :
								]
								SCREAMING_SNAKE_CASE_   :	List[str]													=   in_proj_bias[: config.hidden_size]
								SCREAMING_SNAKE_CASE_   :	Optional[int]													=   in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   in_proj_bias[
								    config.hidden_size : config.hidden_size * 2
								]
								SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   in_proj_weight[
								    -config.hidden_size :, :
								]
								SCREAMING_SNAKE_CASE_   :	str													=   in_proj_bias[-config.hidden_size :]



def 					_lowerCamelCase(			lowerCAmelCase__				:		Optional[Any]      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Dict													=   ['head.weight', 'head.bias']
				for k in ignore_keys:
								state_dict.pop(lowerCAmelCase__    ,    lowerCAmelCase__      )



def 					_lowerCamelCase(			lowerCAmelCase__				:		Any    ,    lowerCAmelCase__				:		List[Any]    ,    lowerCAmelCase__				:		int      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   dct.pop(lowerCAmelCase__      )
				SCREAMING_SNAKE_CASE_   :	Tuple													=   val



@torch.no_grad()
def 					_lowerCamelCase(			lowerCAmelCase__				:		Optional[Any]    ,    lowerCAmelCase__				:		List[Any]      ):





				'''simple docstring'''







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

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

				# Define processor
				SCREAMING_SNAKE_CASE_   :	Tuple													=   ViltImageProcessor(size=384      )
				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   BertTokenizer.from_pretrained('bert-base-uncased'      )
				SCREAMING_SNAKE_CASE_   :	int													=   ViltProcessor(lowerCAmelCase__    ,    lowerCAmelCase__      )

				# Forward pass on example inputs (image + text)
				if nlvr_model:
								SCREAMING_SNAKE_CASE_   :	str													=   Image.open(requests.get('https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'    ,    stream=lowerCAmelCase__      ).raw      )
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   Image.open(requests.get('https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'    ,    stream=lowerCAmelCase__      ).raw      )
								SCREAMING_SNAKE_CASE_   :	int													=   (
								    '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_   :	Optional[int]													=   processor(lowerCAmelCase__    ,    lowerCAmelCase__    ,    return_tensors='pt'      )
								SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   processor(lowerCAmelCase__    ,    lowerCAmelCase__    ,    return_tensors='pt'      )
								SCREAMING_SNAKE_CASE_   :	Any													=   model(
								    input_ids=encoding_a.input_ids    ,    pixel_values=encoding_a.pixel_values    ,    pixel_values_a=encoding_a.pixel_values    ,    )
				else:
								SCREAMING_SNAKE_CASE_   :	List[Any]													=   Image.open(requests.get('http://images.cocodataset.org/val2017/000000039769.jpg'    ,    stream=lowerCAmelCase__      ).raw      )
								if mlm_model:
												SCREAMING_SNAKE_CASE_   :	List[Any]													=   'a bunch of [MASK] laying on a [MASK].'
								else:
												SCREAMING_SNAKE_CASE_   :	str													=   'How many cats are there?'
								SCREAMING_SNAKE_CASE_   :	Optional[Any]													=   processor(lowerCAmelCase__    ,    lowerCAmelCase__    ,    return_tensors='pt'      )
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   model(**lowerCAmelCase__      )

				# Verify outputs
				if mlm_model:
								SCREAMING_SNAKE_CASE_   :	List[Any]													=   torch.Size([1, 11, 3_0522]      )
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   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]    ,    lowerCAmelCase__    ,    atol=1E-4      )

								# verify masked token prediction equals "cats"
								SCREAMING_SNAKE_CASE_   :	Any													=   outputs.logits[0, 4, :].argmax(-1      ).item()
								assert tokenizer.decode([predicted_id]      ) == "cats"
				elif vqa_model:
								SCREAMING_SNAKE_CASE_   :	Dict													=   torch.Size([1, 3129]      )
								SCREAMING_SNAKE_CASE_   :	List[Any]													=   torch.tensor([-15.9_495, -18.1_472, -10.3_041]      )
								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"
								SCREAMING_SNAKE_CASE_   :	Dict													=   outputs.logits.argmax(-1      ).item()
								assert model.config.idalabel[predicted_idx] == "2"
				elif nlvr_model:
								SCREAMING_SNAKE_CASE_   :	str													=   torch.Size([1, 2]      )
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   torch.tensor([-2.8_721, 2.1_291]      )
								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__":
				A        =			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.'
				)

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







import importlib
import inspect
import os
import re


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
A        =			'src/transformers'


# This is to make sure the transformers module imported is the one in the repo.
A        =			importlib.util.spec_from_file_location(
    'transformers',
    os.path.join(PATH_TO_TRANSFORMERS, '__init__.py'),
    submodule_search_locations=[PATH_TO_TRANSFORMERS],
)
A        =			spec.loader.load_module()

A        =			transformers.models.auto.configuration_auto.CONFIG_MAPPING

# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
# For example, `[bert-base-uncased](https://huggingface.co/bert-base-uncased)`
A        =			re.compile('\[(.+?)\]\((https://huggingface\.co/.+?)\)')


A        =			{
    'CLIPConfigMixin',
    'DecisionTransformerConfigMixin',
    'EncoderDecoderConfigMixin',
    'RagConfigMixin',
    'SpeechEncoderDecoderConfigMixin',
    'VisionEncoderDecoderConfigMixin',
    'VisionTextDualEncoderConfigMixin',
}



def 					_lowerCamelCase(			):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_   :	Optional[int]													=   []

				for config_class in list(CONFIG_MAPPING.values()      ):
								SCREAMING_SNAKE_CASE_   :	int													=   False

								# source code of `config_class`
								SCREAMING_SNAKE_CASE_   :	Union[str, Any]													=   inspect.getsource(lowerCAmelCase__      )
								SCREAMING_SNAKE_CASE_   :	List[Any]													=   _re_checkpoint.findall(lowerCAmelCase__      )

								for checkpoint in checkpoints:
												# Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
												# For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
												SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_   :	Tuple													=   checkpoint

												# verify the checkpoint name corresponds to the checkpoint link
												SCREAMING_SNAKE_CASE_   :	List[str]													=   F'''https://huggingface.co/{ckpt_name}'''
												if ckpt_link == ckpt_link_from_name:
																SCREAMING_SNAKE_CASE_   :	Optional[int]													=   True
																break

								SCREAMING_SNAKE_CASE_   :	Tuple													=   config_class.__name__
								if not checkpoint_found and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
												configs_without_checkpoint.append(lowerCAmelCase__      )

				if len(lowerCAmelCase__      ) > 0:
								SCREAMING_SNAKE_CASE_   :	str													=   '\n'.join(sorted(lowerCAmelCase__      )      )
								raise ValueError(F'''The following configurations don\'t contain any valid checkpoint:\n{message}'''      )


if __name__ == "__main__":
				check_config_docstrings_have_checkpoints()