



"""simple docstring"""




import argparse
import collections
import os
import re
import tempfile

import pandas as pd
from datasets import Dataset
from huggingface_hub import hf_hub_download, upload_folder

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/update_metadata.py
__snake_case    :		List[Any]  ='src/transformers'


# This is to make sure the transformers module imported is the one in the repo.
__snake_case    :		Optional[Any]  =direct_transformers_import(TRANSFORMERS_PATH)


# Regexes that match TF/Flax/PT model names.
__snake_case    :		Optional[int]  =re.compile(R'TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
__snake_case    :		Any  =re.compile(R'Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
__snake_case    :		Any  =re.compile(R'(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')


# Fill this with tuples (pipeline_tag, model_mapping, auto_model)
__snake_case    :		Dict  =[
    ('pretraining', 'MODEL_FOR_PRETRAINING_MAPPING_NAMES', 'AutoModelForPreTraining'),
    ('feature-extraction', 'MODEL_MAPPING_NAMES', 'AutoModel'),
    ('audio-classification', 'MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES', 'AutoModelForAudioClassification'),
    ('text-generation', 'MODEL_FOR_CAUSAL_LM_MAPPING_NAMES', 'AutoModelForCausalLM'),
    ('automatic-speech-recognition', 'MODEL_FOR_CTC_MAPPING_NAMES', 'AutoModelForCTC'),
    ('image-classification', 'MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES', 'AutoModelForImageClassification'),
    ('image-segmentation', 'MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES', 'AutoModelForImageSegmentation'),
    ('fill-mask', 'MODEL_FOR_MASKED_LM_MAPPING_NAMES', 'AutoModelForMaskedLM'),
    ('object-detection', 'MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES', 'AutoModelForObjectDetection'),
    (
        'zero-shot-object-detection',
        'MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES',
        'AutoModelForZeroShotObjectDetection',
    ),
    ('question-answering', 'MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES', 'AutoModelForQuestionAnswering'),
    ('text2text-generation', 'MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES', 'AutoModelForSeq2SeqLM'),
    ('text-classification', 'MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES', 'AutoModelForSequenceClassification'),
    ('automatic-speech-recognition', 'MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES', 'AutoModelForSpeechSeq2Seq'),
    (
        'table-question-answering',
        'MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES',
        'AutoModelForTableQuestionAnswering',
    ),
    ('token-classification', 'MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES', 'AutoModelForTokenClassification'),
    ('multiple-choice', 'MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES', 'AutoModelForMultipleChoice'),
    (
        'next-sentence-prediction',
        'MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES',
        'AutoModelForNextSentencePrediction',
    ),
    (
        'audio-frame-classification',
        'MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES',
        'AutoModelForAudioFrameClassification',
    ),
    ('audio-xvector', 'MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES', 'AutoModelForAudioXVector'),
    (
        'document-question-answering',
        'MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES',
        'AutoModelForDocumentQuestionAnswering',
    ),
    (
        'visual-question-answering',
        'MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES',
        'AutoModelForVisualQuestionAnswering',
    ),
    ('image-to-text', 'MODEL_FOR_FOR_VISION_2_SEQ_MAPPING_NAMES', 'AutoModelForVision2Seq'),
    (
        'zero-shot-image-classification',
        'MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES',
        'AutoModelForZeroShotImageClassification',
    ),
    ('depth-estimation', 'MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES', 'AutoModelForDepthEstimation'),
    ('video-classification', 'MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES', 'AutoModelForVideoClassification'),
    ('mask-generation', 'MODEL_FOR_MASK_GENERATION_MAPPING_NAMES', 'AutoModelForMaskGeneration'),
]






def      lowerCAmelCase__		(   lowerCamelCase_		:		Optional[int]):





				'''simple docstring'''

				lowerCAmelCase__					:    Union[str, Any]												=					re.finditer('''.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'''	,__SCREAMING_SNAKE_CASE)
				return [m.group(0) for m in matches]






def      lowerCAmelCase__		(   ):





				'''simple docstring'''

				lowerCAmelCase__					:    Optional[int]												=					transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
				lowerCAmelCase__					:    Tuple												=					{
				    config.replace('''Config'''	,''''''): model_type for model_type, config in config_maping_names.items()
				}

				# Dictionaries flagging if each model prefix has a backend in PT/TF/Flax.
				lowerCAmelCase__					:    str												=					collections.defaultdict(__SCREAMING_SNAKE_CASE)
				lowerCAmelCase__					:    Tuple												=					collections.defaultdict(__SCREAMING_SNAKE_CASE)
				lowerCAmelCase__					:    Union[str, Any]												=					collections.defaultdict(__SCREAMING_SNAKE_CASE)

				# Let's lookup through all transformers object (once) and find if models are supported by a given backend.
				for attr_name in dir(__SCREAMING_SNAKE_CASE):
								lowerCAmelCase__					:    str												=					None
								if _re_tf_models.match(__SCREAMING_SNAKE_CASE) is not None:
												lowerCAmelCase__					:    List[Any]												=					tf_models
												lowerCAmelCase__					:    Dict												=					_re_tf_models.match(__SCREAMING_SNAKE_CASE).groups()[0]
								elif _re_flax_models.match(__SCREAMING_SNAKE_CASE) is not None:
												lowerCAmelCase__					:    List[str]												=					flax_models
												lowerCAmelCase__					:    Any												=					_re_flax_models.match(__SCREAMING_SNAKE_CASE).groups()[0]
								elif _re_pt_models.match(__SCREAMING_SNAKE_CASE) is not None:
												lowerCAmelCase__					:    Union[str, Any]												=					pt_models
												lowerCAmelCase__					:    Union[str, Any]												=					_re_pt_models.match(__SCREAMING_SNAKE_CASE).groups()[0]

								if lookup_dict is not None:
												while len(__SCREAMING_SNAKE_CASE) > 0:
																if attr_name in model_prefix_to_model_type:
																				lowerCAmelCase__					:    Any												=					True
																				break
																# Try again after removing the last word in the name
																lowerCAmelCase__					:    int												=					"""""".join(camel_case_split(__SCREAMING_SNAKE_CASE)[:-1])

				lowerCAmelCase__					:    Union[str, Any]												=					set(list(pt_models.keys()) + list(tf_models.keys()) + list(flax_models.keys()))
				lowerCAmelCase__					:    List[Any]												=					list(__SCREAMING_SNAKE_CASE)
				all_models.sort()

				lowerCAmelCase__					:    Dict												=					{"""model_type""": all_models}
				lowerCAmelCase__					:    Optional[Any]												=					[pt_models[t] for t in all_models]
				lowerCAmelCase__					:    List[Any]												=					[tf_models[t] for t in all_models]
				lowerCAmelCase__					:    Tuple												=					[flax_models[t] for t in all_models]

				# Now let's use the auto-mapping names to make sure
				lowerCAmelCase__					:    Any												=					{}
				for t in all_models:
								if t in transformers_module.models.auto.processing_auto.PROCESSOR_MAPPING_NAMES:
												lowerCAmelCase__					:    Union[str, Any]												=					"""AutoProcessor"""
								elif t in transformers_module.models.auto.tokenization_auto.TOKENIZER_MAPPING_NAMES:
												lowerCAmelCase__					:    Dict												=					"""AutoTokenizer"""
								elif t in transformers_module.models.auto.feature_extraction_auto.FEATURE_EXTRACTOR_MAPPING_NAMES:
												lowerCAmelCase__					:    Tuple												=					"""AutoFeatureExtractor"""
								else:
												# Default to AutoTokenizer if a model has nothing, for backward compatibility.
												lowerCAmelCase__					:    Tuple												=					"""AutoTokenizer"""

				lowerCAmelCase__					:    Optional[int]												=					[processors[t] for t in all_models]

				return pd.DataFrame(__SCREAMING_SNAKE_CASE)






def      lowerCAmelCase__		(   lowerCamelCase_		:		List[Any]):





				'''simple docstring'''

				lowerCAmelCase__					:    Union[str, Any]												=					[
				    transformers_module.models.auto.modeling_auto,
				    transformers_module.models.auto.modeling_tf_auto,
				    transformers_module.models.auto.modeling_flax_auto,
				]
				for pipeline_tag, model_mapping, auto_class in PIPELINE_TAGS_AND_AUTO_MODELS:
								lowerCAmelCase__					:    str												=					[model_mapping, f"""TF_{model_mapping}""", f"""FLAX_{model_mapping}"""]
								lowerCAmelCase__					:    int												=					[auto_class, f"""TF_{auto_class}""", f"""Flax_{auto_class}"""]
								# Loop through all three frameworks
								for module, cls, mapping in zip(__SCREAMING_SNAKE_CASE	,__SCREAMING_SNAKE_CASE	,__SCREAMING_SNAKE_CASE):
												# The type of pipeline may not exist in this framework
												if not hasattr(__SCREAMING_SNAKE_CASE	,__SCREAMING_SNAKE_CASE):
																continue
												# First extract all model_names
												lowerCAmelCase__					:    List[str]												=					[]
												for name in getattr(__SCREAMING_SNAKE_CASE	,__SCREAMING_SNAKE_CASE).values():
																if isinstance(__SCREAMING_SNAKE_CASE	,__SCREAMING_SNAKE_CASE):
																				model_names.append(__SCREAMING_SNAKE_CASE)
																else:
																				model_names.extend(list(__SCREAMING_SNAKE_CASE))

            # Add pipeline tag and auto model class for those models
												table.update({model_name: (pipeline_tag, cls) for model_name in model_names})

				return table






def      lowerCAmelCase__		(   lowerCamelCase_		:		str	,lowerCamelCase_		:		List[Any]):





				'''simple docstring'''

				lowerCAmelCase__					:    Dict												=					get_frameworks_table()
				lowerCAmelCase__					:    Tuple												=					Dataset.from_pandas(__SCREAMING_SNAKE_CASE)

				lowerCAmelCase__					:    int												=					hf_hub_download(
				    '''huggingface/transformers-metadata'''	,'''pipeline_tags.json'''	,repo_type='''dataset'''	,token=__SCREAMING_SNAKE_CASE)
				lowerCAmelCase__					:    Dict												=					Dataset.from_json(__SCREAMING_SNAKE_CASE)
				lowerCAmelCase__					:    Any												=					{
				    tags_dataset[i]["""model_class"""]: (tags_dataset[i]["""pipeline_tag"""], tags_dataset[i]["""auto_class"""])
				    for i in range(len(__SCREAMING_SNAKE_CASE))
				}
				lowerCAmelCase__					:    Optional[int]												=					update_pipeline_and_auto_class_table(__SCREAMING_SNAKE_CASE)

				# Sort the model classes to avoid some nondeterministic updates to create false update commits.
				lowerCAmelCase__					:    Optional[int]												=					sorted(table.keys())
				lowerCAmelCase__					:    Union[str, Any]												=					pd.DataFrame(
				    {
				        '''model_class''': model_classes,
				        '''pipeline_tag''': [table[m][0] for m in model_classes],
				        '''auto_class''': [table[m][1] for m in model_classes],
				    })
				lowerCAmelCase__					:    Any												=					Dataset.from_pandas(__SCREAMING_SNAKE_CASE)

				with tempfile.TemporaryDirectory() as tmp_dir:
								frameworks_dataset.to_json(os.path.join(__SCREAMING_SNAKE_CASE	,'''frameworks.json'''))
								tags_dataset.to_json(os.path.join(__SCREAMING_SNAKE_CASE	,'''pipeline_tags.json'''))

								if commit_sha is not None:
												lowerCAmelCase__					:    Dict												=					(
												    f"""Update with commit {commit_sha}\n\nSee: """
												    f"""https://github.com/huggingface/transformers/commit/{commit_sha}"""
												)
								else:
												lowerCAmelCase__					:    List[Any]												=					"""Update"""

								upload_folder(
								    repo_id='''huggingface/transformers-metadata'''	,folder_path=__SCREAMING_SNAKE_CASE	,repo_type='''dataset'''	,token=__SCREAMING_SNAKE_CASE	,commit_message=__SCREAMING_SNAKE_CASE	,)






def      lowerCAmelCase__		(   ):





				'''simple docstring'''

				lowerCAmelCase__					:    List[Any]												=					{tag: cls for tag, _, cls in PIPELINE_TAGS_AND_AUTO_MODELS}
				lowerCAmelCase__					:    Dict												=					transformers_module.pipelines.SUPPORTED_TASKS
				lowerCAmelCase__					:    Tuple												=					[]
				for key in pipeline_tasks:
								if key not in in_table:
												lowerCAmelCase__					:    Dict												=					pipeline_tasks[key]["""pt"""]
												if isinstance(__SCREAMING_SNAKE_CASE	,(list, tuple)):
																lowerCAmelCase__					:    Tuple												=					model[0]
												lowerCAmelCase__					:    Optional[Any]												=					model.__name__
												if model not in in_table.values():
																missing.append(__SCREAMING_SNAKE_CASE)

				if len(__SCREAMING_SNAKE_CASE) > 0:
								lowerCAmelCase__					:    Union[str, Any]												=					""", """.join(__SCREAMING_SNAKE_CASE)
								raise ValueError(
								    '''The following pipeline tags are not present in the `PIPELINE_TAGS_AND_AUTO_MODELS` constant inside '''
								    f"""`utils/update_metadata.py`: {msg}. Please add them!""")


if __name__ == "__main__":
							__snake_case    :		List[str]  =argparse.ArgumentParser()
							parser.add_argument('--token', type=str, help='The token to use to push to the transformers-metadata dataset.')
							parser.add_argument('--commit_sha', type=str, help='The sha of the commit going with this update.')
							parser.add_argument('--check-only', action='store_true', help='Activate to just check all pipelines are present.')
							__snake_case    :		Any  =parser.parse_args()

							if args.check_only:
														check_pipeline_tags()
							else:
														update_metadata(args.token, args.commit_sha)






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()
__snake_case    :		Tuple  =logging.get_logger(__name__)






def      lowerCAmelCase__		(   lowerCamelCase_		:		Tuple	,lowerCamelCase_		:		Dict=False	,lowerCamelCase_		:		List[Any]=False	,lowerCamelCase_		:		int=False):





				'''simple docstring'''

				lowerCAmelCase__					:    Tuple												=					[]
				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_		:		Dict	,lowerCamelCase_		:		Optional[Any]):





				'''simple docstring'''

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






def      lowerCAmelCase__		(   lowerCamelCase_		:		List[str]):





				'''simple docstring'''

				lowerCAmelCase__					:    List[Any]												=					['''head.weight''', '''head.bias''']
				for k in ignore_keys:
								state_dict.pop(lowerCamelCase_	,lowerCamelCase_)






def      lowerCAmelCase__		(   lowerCamelCase_		:		Optional[Any]	,lowerCamelCase_		:		Optional[int]	,lowerCamelCase_		:		int):





				'''simple docstring'''

				lowerCAmelCase__					:    int												=					dct.pop(lowerCamelCase_)
				lowerCAmelCase__					:    List[Any]												=					val






@torch.no_grad()
def      lowerCAmelCase__		(   lowerCamelCase_		:		int	,lowerCamelCase_		:		Union[str, Any]):





				'''simple docstring'''

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

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

				# Define processor
				lowerCAmelCase__					:    List[str]												=					ViltImageProcessor(size=384)
				lowerCAmelCase__					:    Tuple												=					BertTokenizer.from_pretrained('''bert-base-uncased''')
				lowerCAmelCase__					:    Union[str, Any]												=					ViltProcessor(lowerCamelCase_	,lowerCamelCase_)

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

				# Verify outputs
				if mlm_model:
								lowerCAmelCase__					:    Dict												=					torch.Size([1, 11, 30522])
								lowerCAmelCase__					:    int												=					torch.tensor([-12.5061, -12.5123, -12.5174])
								assert outputs.logits.shape == expected_shape
								assert torch.allclose(outputs.logits[0, 0, :3]	,lowerCamelCase_	,atol=1E-4)

								# verify masked token prediction equals "cats"
								lowerCAmelCase__					:    Optional[Any]												=					outputs.logits[0, 4, :].argmax(-1).item()
								assert tokenizer.decode([predicted_id]) == "cats"
				elif vqa_model:
								lowerCAmelCase__					:    List[Any]												=					torch.Size([1, 3129])
								lowerCAmelCase__					:    str												=					torch.tensor([-15.9495, -18.1472, -10.3041])
								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"
								lowerCAmelCase__					:    List[Any]												=					outputs.logits.argmax(-1).item()
								assert model.config.idalabel[predicted_idx] == "2"
				elif nlvr_model:
								lowerCAmelCase__					:    Union[str, Any]												=					torch.Size([1, 2])
								lowerCAmelCase__					:    Dict												=					torch.tensor([-2.8721, 2.1291])
								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__":
							__snake_case    :		Dict  =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.'
							)

							__snake_case    :		Union[str, Any]  =parser.parse_args()
							convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)

