





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 						lowerCamelCase__     (     UpperCamelCase__			:       Dict       ,    UpperCamelCase__			:       str=False       ,    UpperCamelCase__			:       int=False       ,    UpperCamelCase__			:       Any=False						)			-> int:



					'''simple docstring'''




					_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__     (     UpperCamelCase__			:       int       ,    UpperCamelCase__			:       int						)			-> Optional[Any]:



					'''simple docstring'''




					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__     (     UpperCamelCase__			:       int						)			-> Optional[Any]:



					'''simple docstring'''




					_snake_case									=				['head.weight', 'head.bias']
					for k in ignore_keys:
										state_dict.pop(UpperCamelCase__       ,    UpperCamelCase__						)





def 						lowerCamelCase__     (     UpperCamelCase__			:       Optional[Any]       ,    UpperCamelCase__			:       List[str]       ,    UpperCamelCase__			:       Dict						)			-> Any:



					'''simple docstring'''




					_snake_case									=				dct.pop(UpperCamelCase__						)
					_snake_case									=				val





@torch.no_grad()
def 						lowerCamelCase__     (     UpperCamelCase__			:       str       ,    UpperCamelCase__			:       List[Any]						)			-> Any:



					'''simple docstring'''




					_snake_case									=				ViltConfig(image_size=384       ,    patch_size=32       ,    tie_word_embeddings=UpperCamelCase__						)
					_snake_case									=				False
					_snake_case									=				False
					_snake_case									=				False
					_snake_case									=				False
					if "vqa" in checkpoint_url:
										_snake_case									=				True
										_snake_case									=				3_129
										_snake_case									=				'huggingface/label-files'
										_snake_case									=				'vqa2-id2label.json'
										_snake_case									=				json.load(open(hf_hub_download(UpperCamelCase__       ,    UpperCamelCase__       ,    repo_type='dataset'						)       ,    'r'						)						)
										_snake_case									=				{int(UpperCamelCase__						): v for k, v in idalabel.items()}
										_snake_case									=				idalabel
										_snake_case									=				{v: k for k, v in idalabel.items()}
										_snake_case									=				ViltForQuestionAnswering(UpperCamelCase__						)
					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(UpperCamelCase__						)
					elif "irtr" in checkpoint_url:
										_snake_case									=				True
										_snake_case									=				ViltForImageAndTextRetrieval(UpperCamelCase__						)
					elif "mlm_itm" in checkpoint_url:
										_snake_case									=				True
										_snake_case									=				ViltForMaskedLM(UpperCamelCase__						)
					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(UpperCamelCase__       ,    map_location='cpu'						)['state_dict']
					_snake_case									=				create_rename_keys(UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__						)
					for src, dest in rename_keys:
										rename_key(UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__						)
					read_in_q_k_v(UpperCamelCase__       ,    UpperCamelCase__						)
					if mlm_model or irtr_model:
										_snake_case									=				['itm_score.fc.weight', 'itm_score.fc.bias']
										for k in ignore_keys:
															state_dict.pop(UpperCamelCase__       ,    UpperCamelCase__						)

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

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

					# 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=UpperCamelCase__						).raw						)
										_snake_case									=				Image.open(requests.get('https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'       ,    stream=UpperCamelCase__						).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(UpperCamelCase__       ,    UpperCamelCase__       ,    return_tensors='pt'						)
										_snake_case									=				processor(UpperCamelCase__       ,    UpperCamelCase__       ,    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=UpperCamelCase__						).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(UpperCamelCase__       ,    UpperCamelCase__       ,    return_tensors='pt'						)
										_snake_case									=				model(**UpperCamelCase__						)

					# Verify outputs
					if mlm_model:
										_snake_case									=				torch.Size([1, 11, 30_522]						)
										_snake_case									=				torch.tensor([-12.5061, -12.5123, -12.5174]						)
										assert outputs.logits.shape == expected_shape
										assert torch.allclose(outputs.logits[0, 0, :3]       ,    UpperCamelCase__       ,    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, 3_129]						)
										_snake_case									=				torch.tensor([-15.9495, -18.1472, -10.3041]						)
										assert torch.allclose(outputs.logits[0, :3]       ,    UpperCamelCase__       ,    atol=1e-4						)
										assert outputs.logits.shape == expected_shape
										assert torch.allclose(outputs.logits[0, 0, :3]       ,    UpperCamelCase__       ,    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]       ,    UpperCamelCase__       ,    atol=1e-4						)
										assert outputs.logits.shape == expected_shape

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


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)










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase_      							=     {"""configuration_swin""": ["""SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP""", """SwinConfig""", """SwinOnnxConfig"""]}


try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase_      							=     [
					    """SWIN_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """SwinForImageClassification""",
					    """SwinForMaskedImageModeling""",
					    """SwinModel""",
					    """SwinPreTrainedModel""",
					    """SwinBackbone""",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase_      							=     [
					    """TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """TFSwinForImageClassification""",
					    """TFSwinForMaskedImageModeling""",
					    """TFSwinModel""",
					    """TFSwinPreTrainedModel""",
					]

if TYPE_CHECKING:
					from .configuration_swin import SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinConfig, SwinOnnxConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_swin import (
										    SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
										    SwinBackbone,
										    SwinForImageClassification,
										    SwinForMaskedImageModeling,
										    SwinModel,
										    SwinPreTrainedModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_swin import (
										    TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFSwinForImageClassification,
										    TFSwinForMaskedImageModeling,
										    TFSwinModel,
										    TFSwinPreTrainedModel,
										)

else:
					import sys

					UpperCAmelCase_      							=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



