




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 __UpperCamelCase			(       _A ,  _A=False ,  _A=False ,  _A=False  ):
						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 __UpperCamelCase			(       _A ,  _A  ):
						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 __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  ['''head.weight''', '''head.bias''']
						for k in ignore_keys:
												state_dict.pop(_A ,  _A  )





def __UpperCamelCase			(       _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  dct.pop(_A  )
						lowerCAmelCase_      			=  val





@torch.no_grad()
def __UpperCamelCase			(       _A ,  _A  ):
						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([-1_2.5_0_6_1, -1_2.5_1_2_3, -1_2.5_1_7_4]  )
												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([-1_5.9_4_9_5, -1_8.1_4_7_2, -1_0.3_0_4_1]  )
												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_7_2_1, 2.1_2_9_1]  )
												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__":
						_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 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 (
    SwiftFormerConfig,
    SwiftFormerForImageClassification,
    ViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
_A     =			logging.get_logger(__name__)

_A     =			torch.device('''cpu''')





def __UpperCamelCase			(       ):
						lowerCAmelCase_      			=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
						lowerCAmelCase_      			=  Image.open(requests.get(_A ,  stream=_A  ).raw  )
						return im





def __UpperCamelCase			(       _A  ):
						if swiftformer_name == "swiftformer_xs":
												return torch.tensor([-2.1_7_0_3E0_0, 2.1_1_0_7E0_0, -2.0_8_1_1E0_0, 8.8_6_8_5E-0_1, 2.4_3_6_0E-0_1]  )

						elif swiftformer_name == "swiftformer_s":
												return torch.tensor([3.9_6_3_6E-0_1, 2.3_4_7_8E-0_1, -1.6_9_6_3E0_0, -1.7_3_8_1E0_0, -8.6_3_3_7E-0_1]  )

						elif swiftformer_name == "swiftformer_l1":
												return torch.tensor([-4.2_7_6_8E-0_1, -4.7_4_2_9E-0_1, -1.0_8_9_7E0_0, -1.0_2_4_8E0_0, 3.5_5_2_3E-0_2]  )

						elif swiftformer_name == "swiftformer_l3":
												return torch.tensor([-2.5_3_3_0E-0_1, 2.4_2_1_1E-0_1, -6.0_1_8_5E-0_1, -8.2_7_8_9E-0_1, -6.0_4_4_6E-0_2]  )





def __UpperCamelCase			(       _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  dct.pop(_A  )
						lowerCAmelCase_      			=  val





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  []
						for k in state_dict.keys():
												lowerCAmelCase_      			=  k
												if ".pwconv" in k:
																		lowerCAmelCase_      			=  k_new.replace('''.pwconv''' ,  '''.point_wise_conv'''  )
												if ".dwconv" in k:
																		lowerCAmelCase_      			=  k_new.replace('''.dwconv''' ,  '''.depth_wise_conv'''  )
												if ".Proj." in k:
																		lowerCAmelCase_      			=  k_new.replace('''.Proj.''' ,  '''.proj.'''  )
												if "patch_embed" in k_new:
																		lowerCAmelCase_      			=  k_new.replace('''patch_embed''' ,  '''swiftformer.patch_embed.patch_embedding'''  )
												if "network" in k_new:
																		lowerCAmelCase_      			=  k_new.split('''.'''  )
																		if ls[2].isdigit():
																								lowerCAmelCase_      			=  '''swiftformer.encoder.network.''' + ls[1] + '''.blocks.''' + ls[2] + '''.''' + '''.'''.join(ls[3:]  )
																		else:
																								lowerCAmelCase_      			=  k_new.replace('''network''' ,  '''swiftformer.encoder.network'''  )
												rename_keys.append((k, k_new)  )
						return rename_keys





@torch.no_grad()
def __UpperCamelCase			(       _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  SwiftFormerConfig()

						# dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
						lowerCAmelCase_      			=  1000
						lowerCAmelCase_      			=  '''huggingface/label-files'''
						lowerCAmelCase_      			=  '''imagenet-1k-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()}

						# size of the architecture
						if swiftformer_name == "swiftformer_xs":
												lowerCAmelCase_      			=  [3, 3, 6, 4]
												lowerCAmelCase_      			=  [48, 56, 112, 220]

						elif swiftformer_name == "swiftformer_s":
												lowerCAmelCase_      			=  [3, 3, 9, 6]
												lowerCAmelCase_      			=  [48, 64, 168, 224]

						elif swiftformer_name == "swiftformer_l1":
												lowerCAmelCase_      			=  [4, 3, 10, 5]
												lowerCAmelCase_      			=  [48, 96, 192, 384]

						elif swiftformer_name == "swiftformer_l3":
												lowerCAmelCase_      			=  [4, 4, 12, 6]
												lowerCAmelCase_      			=  [64, 128, 320, 512]

						# load state_dict of original model, remove and rename some keys
						if original_ckpt:
												if original_ckpt.startswith('''https'''  ):
																		lowerCAmelCase_      			=  torch.hub.load_state_dict_from_url(_A ,  map_location='''cpu''' ,  check_hash=_A  )
												else:
																		lowerCAmelCase_      			=  torch.load(_A ,  map_location='''cpu'''  )
						lowerCAmelCase_      			=  checkpoint

						lowerCAmelCase_      			=  create_rename_keys(_A  )
						for rename_key_src, rename_key_dest in rename_keys:
												rename_key(_A ,  _A ,  _A  )

						# load HuggingFace model
						lowerCAmelCase_      			=  SwiftFormerForImageClassification(_A  ).eval()
						hf_model.load_state_dict(_A  )

						# prepare test inputs
						lowerCAmelCase_      			=  prepare_img()
						lowerCAmelCase_      			=  ViTImageProcessor.from_pretrained('''preprocessor_config'''  )
						lowerCAmelCase_      			=  processor(images=_A ,  return_tensors='''pt'''  )

						# compare outputs from both models
						lowerCAmelCase_      			=  get_expected_output(_A  )
						lowerCAmelCase_      			=  hf_model(inputs['''pixel_values''']  ).logits

						assert hf_logits.shape == torch.Size([1, 1000]  )
						assert torch.allclose(hf_logits[0, 0:5] ,  _A ,  atol=1E-3  )

						Path(_A  ).mkdir(exist_ok=_A  )
						print(f"Saving model {swiftformer_name} to {pytorch_dump_folder_path}"  )
						hf_model.save_pretrained(_A  )


if __name__ == "__main__":
						_A     =			argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '''--swiftformer_name''',
						    default='''swiftformer_xs''',
						    choices=['''swiftformer_xs''', '''swiftformer_s''', '''swiftformer_l1''', '''swiftformer_l3'''],
						    type=str,
						    help='''Name of the SwiftFormer model you\'d like to convert.''',
						)
						parser.add_argument(
						    '''--pytorch_dump_folder_path''',
						    default='''./converted_outputs/''',
						    type=str,
						    help='''Path to the output PyTorch model directory.''',
						)
						parser.add_argument('''--original_ckpt''', default=None, type=str, help='''Path to the original model checkpoint.''')

						_A     =			parser.parse_args()
						convert_swiftformer_checkpoint(args.swiftformer_name, args.pytorch_dump_folder_path, args.original_ckpt)




