


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







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







def     _UpperCamelCase      (		UpperCamelCase__					):
			UpperCAmelCase__  :					Optional[int]        =     ["""head.weight""", """head.bias"""]
			for k in ignore_keys:
						state_dict.pop(UpperCamelCase__ ,      UpperCamelCase__					)







def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__					):
			UpperCAmelCase__  :					Optional[int]        =     dct.pop(UpperCamelCase__					)
			UpperCAmelCase__  :					List[str]        =     val







@torch.no_grad()
def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__					):
			UpperCAmelCase__  :					str        =     ViltConfig(image_size=3_8_4 ,      patch_size=3_2 ,      tie_word_embeddings=UpperCamelCase__					)
			UpperCAmelCase__  :					Union[str, Any]        =     False
			UpperCAmelCase__  :					List[Any]        =     False
			UpperCAmelCase__  :					Any        =     False
			UpperCAmelCase__  :					int        =     False
			if "vqa" in checkpoint_url:
						UpperCAmelCase__  :					Union[str, Any]        =     True
						UpperCAmelCase__  :					Optional[Any]        =     3_1_2_9
						UpperCAmelCase__  :					Optional[Any]        =     """huggingface/label-files"""
						UpperCAmelCase__  :					int        =     """vqa2-id2label.json"""
						UpperCAmelCase__  :					Optional[Any]        =     json.load(open(hf_hub_download(UpperCamelCase__ ,      UpperCamelCase__ ,      repo_type="""dataset"""					) ,      """r"""					)					)
						UpperCAmelCase__  :					int        =     {int(UpperCamelCase__					): v for k, v in idalabel.items()}
						UpperCAmelCase__  :					List[str]        =     idalabel
						UpperCAmelCase__  :					Any        =     {v: k for k, v in idalabel.items()}
						UpperCAmelCase__  :					Optional[Any]        =     ViltForQuestionAnswering(UpperCamelCase__					)
			elif "nlvr" in checkpoint_url:
						UpperCAmelCase__  :					Union[str, Any]        =     True
						UpperCAmelCase__  :					int        =     2
						UpperCAmelCase__  :					List[Any]        =     {0: """False""", 1: """True"""}
						UpperCAmelCase__  :					Union[str, Any]        =     {v: k for k, v in config.idalabel.items()}
						UpperCAmelCase__  :					str        =     3
						UpperCAmelCase__  :					Any        =     ViltForImagesAndTextClassification(UpperCamelCase__					)
			elif "irtr" in checkpoint_url:
						UpperCAmelCase__  :					Any        =     True
						UpperCAmelCase__  :					Dict        =     ViltForImageAndTextRetrieval(UpperCamelCase__					)
			elif "mlm_itm" in checkpoint_url:
						UpperCAmelCase__  :					Dict        =     True
						UpperCAmelCase__  :					List[str]        =     ViltForMaskedLM(UpperCamelCase__					)
			else:
						raise ValueError("""Unknown model type"""					)

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

			# Define processor
			UpperCAmelCase__  :					List[Any]        =     ViltImageProcessor(size=3_8_4					)
			UpperCAmelCase__  :					str        =     BertTokenizer.from_pretrained("""bert-base-uncased"""					)
			UpperCAmelCase__  :					int        =     ViltProcessor(UpperCamelCase__ ,      UpperCamelCase__					)

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

			# Verify outputs
			if mlm_model:
						UpperCAmelCase__  :					List[Any]        =     torch.Size([1, 1_1, 3_0_5_2_2]					)
						UpperCAmelCase__  :					List[Any]        =     torch.tensor([-12.50_61, -12.51_23, -12.51_74]					)
						assert outputs.logits.shape == expected_shape
						assert torch.allclose(outputs.logits[0, 0, :3] ,      UpperCamelCase__ ,      atol=1e-4					)

						# verify masked token prediction equals "cats"
						UpperCAmelCase__  :					Any        =     outputs.logits[0, 4, :].argmax(-1					).item()
						assert tokenizer.decode([predicted_id]					) == "cats"
			elif vqa_model:
						UpperCAmelCase__  :					Optional[int]        =     torch.Size([1, 3_1_2_9]					)
						UpperCAmelCase__  :					str        =     torch.tensor([-15.94_95, -18.14_72, -10.30_41]					)
						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"
						UpperCAmelCase__  :					Union[str, Any]        =     outputs.logits.argmax(-1					).item()
						assert model.config.idalabel[predicted_idx] == "2"
			elif nlvr_model:
						UpperCAmelCase__  :					Any        =     torch.Size([1, 2]					)
						UpperCAmelCase__  :					Optional[Any]        =     torch.tensor([-2.87_21, 2.12_91]					)
						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__":
		__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)



'''simple docstring'''



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__A	=logging.get_logger(__name__)

__A	={
    'xlm-mlm-en-2048': 'https://huggingface.co/xlm-mlm-en-2048/resolve/main/config.json',
    'xlm-mlm-ende-1024': 'https://huggingface.co/xlm-mlm-ende-1024/resolve/main/config.json',
    'xlm-mlm-enfr-1024': 'https://huggingface.co/xlm-mlm-enfr-1024/resolve/main/config.json',
    'xlm-mlm-enro-1024': 'https://huggingface.co/xlm-mlm-enro-1024/resolve/main/config.json',
    'xlm-mlm-tlm-xnli15-1024': 'https://huggingface.co/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json',
    'xlm-mlm-xnli15-1024': 'https://huggingface.co/xlm-mlm-xnli15-1024/resolve/main/config.json',
    'xlm-clm-enfr-1024': 'https://huggingface.co/xlm-clm-enfr-1024/resolve/main/config.json',
    'xlm-clm-ende-1024': 'https://huggingface.co/xlm-clm-ende-1024/resolve/main/config.json',
    'xlm-mlm-17-1280': 'https://huggingface.co/xlm-mlm-17-1280/resolve/main/config.json',
    'xlm-mlm-100-1280': 'https://huggingface.co/xlm-mlm-100-1280/resolve/main/config.json',
}


class 		_snake_case					(     a__	):
					lowerCAmelCase      :Any   =					'''xlm'''
					lowerCAmelCase      :Any   =					{
					    '''hidden_size''': '''emb_dim''',
					    '''num_attention_heads''': '''n_heads''',
					    '''num_hidden_layers''': '''n_layers''',
					    '''n_words''': '''vocab_size''',  # For backward compatibility
					}




					def __init__(      self							,							_lowerCamelCase=3_0145							,							_lowerCamelCase=2048							,							_lowerCamelCase=12							,							_lowerCamelCase=16							,							_lowerCamelCase=0.1							,							_lowerCamelCase=0.1							,							_lowerCamelCase=True							,							_lowerCamelCase=False							,							_lowerCamelCase=False							,							_lowerCamelCase=False							,							_lowerCamelCase=1							,							_lowerCamelCase=True							,							_lowerCamelCase=512							,							_lowerCamelCase=2048**-0.5							,							_lowerCamelCase=1e-1_2							,							_lowerCamelCase=0.02							,							_lowerCamelCase=0							,							_lowerCamelCase=1							,							_lowerCamelCase=2							,							_lowerCamelCase=3							,							_lowerCamelCase=5							,							_lowerCamelCase=True							,							_lowerCamelCase="first"							,							_lowerCamelCase=True							,							_lowerCamelCase=None							,							_lowerCamelCase=True							,							_lowerCamelCase=0.1							,							_lowerCamelCase=5							,							_lowerCamelCase=5							,							_lowerCamelCase=0							,							_lowerCamelCase=0							,							_lowerCamelCase=2							,							_lowerCamelCase=0							,							**_lowerCamelCase							,							):
								UpperCAmelCase__  :					Tuple        =     vocab_size
								UpperCAmelCase__  :					Tuple        =     emb_dim
								UpperCAmelCase__  :					Optional[Any]        =     n_layers
								UpperCAmelCase__  :					List[str]        =     n_heads
								UpperCAmelCase__  :					List[str]        =     dropout
								UpperCAmelCase__  :					Optional[int]        =     attention_dropout
								UpperCAmelCase__  :					Tuple        =     gelu_activation
								UpperCAmelCase__  :					Optional[Any]        =     sinusoidal_embeddings
								UpperCAmelCase__  :					int        =     causal
								UpperCAmelCase__  :					Union[str, Any]        =     asm
								UpperCAmelCase__  :					Optional[Any]        =     n_langs
								UpperCAmelCase__  :					List[Any]        =     use_lang_emb
								UpperCAmelCase__  :					Union[str, Any]        =     layer_norm_eps
								UpperCAmelCase__  :					List[str]        =     bos_index
								UpperCAmelCase__  :					List[Any]        =     eos_index
								UpperCAmelCase__  :					int        =     pad_index
								UpperCAmelCase__  :					str        =     unk_index
								UpperCAmelCase__  :					Dict        =     mask_index
								UpperCAmelCase__  :					str        =     is_encoder
								UpperCAmelCase__  :					Dict        =     max_position_embeddings
								UpperCAmelCase__  :					Any        =     embed_init_std
								UpperCAmelCase__  :					List[Any]        =     init_std
								UpperCAmelCase__  :					List[str]        =     summary_type
								UpperCAmelCase__  :					Union[str, Any]        =     summary_use_proj
								UpperCAmelCase__  :					Any        =     summary_activation
								UpperCAmelCase__  :					List[str]        =     summary_proj_to_labels
								UpperCAmelCase__  :					Union[str, Any]        =     summary_first_dropout
								UpperCAmelCase__  :					str        =     start_n_top
								UpperCAmelCase__  :					str        =     end_n_top
								UpperCAmelCase__  :					Tuple        =     mask_token_id
								UpperCAmelCase__  :					Union[str, Any]        =     lang_id

								if "n_words" in kwargs:
											UpperCAmelCase__  :					List[str]        =     kwargs["""n_words"""]

								super().__init__(pad_token_id=_lowerCamelCase							,							bos_token_id=_lowerCamelCase							,							**_lowerCamelCase)


class 		_snake_case					(     a__	):




					@property
					def 				snake_case__						(      self):
								if self.task == "multiple-choice":
											UpperCAmelCase__  :					Any        =     {0: """batch""", 1: """choice""", 2: """sequence"""}
								else:
											UpperCAmelCase__  :					List[Any]        =     {0: """batch""", 1: """sequence"""}
								return OrderedDict(
								    [
								        ("""input_ids""", dynamic_axis),
								        ("""attention_mask""", dynamic_axis),
								        ("""token_type_ids""", dynamic_axis),
								    ])