


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


def    lowerCamelCase__   ( A     :    Optional[int]   ,     A     :    int=False   ,     A     :    Dict=False   ,     A     :    Optional[int]=False  ):




						'''simple docstring'''






						UpperCAmelCase         =			[]
						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__   ( A     :    int   ,     A     :    Tuple  ):




						'''simple docstring'''






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


def    lowerCamelCase__   ( A     :    Optional[Any]  ):




						'''simple docstring'''






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


def    lowerCamelCase__   ( A     :    Optional[Any]   ,     A     :    Optional[Any]   ,     A     :    Dict  ):




						'''simple docstring'''






						UpperCAmelCase         =			dct.pop(A  )
						UpperCAmelCase         =			val


@torch.no_grad()
def    lowerCamelCase__   ( A     :    List[str]   ,     A     :    Dict  ):




						'''simple docstring'''






						UpperCAmelCase         =			ViltConfig(image_size=3_84   ,     patch_size=32   ,     tie_word_embeddings=A  )
						UpperCAmelCase         =			False
						UpperCAmelCase         =			False
						UpperCAmelCase         =			False
						UpperCAmelCase         =			False
						if "vqa" in checkpoint_url:
												UpperCAmelCase         =			True
												UpperCAmelCase         =			31_29
												UpperCAmelCase         =			'''huggingface/label-files'''
												UpperCAmelCase         =			'''vqa2-id2label.json'''
												UpperCAmelCase         =			json.load(open(hf_hub_download(A   ,     A   ,     repo_type='''dataset'''  )   ,     '''r'''  )  )
												UpperCAmelCase         =			{int(A  ): v for k, v in idalabel.items()}
												UpperCAmelCase         =			idalabel
												UpperCAmelCase         =			{v: k for k, v in idalabel.items()}
												UpperCAmelCase         =			ViltForQuestionAnswering(A  )
						elif "nlvr" in checkpoint_url:
												UpperCAmelCase         =			True
												UpperCAmelCase         =			2
												UpperCAmelCase         =			{0: '''False''', 1: '''True'''}
												UpperCAmelCase         =			{v: k for k, v in config.idalabel.items()}
												UpperCAmelCase         =			3
												UpperCAmelCase         =			ViltForImagesAndTextClassification(A  )
						elif "irtr" in checkpoint_url:
												UpperCAmelCase         =			True
												UpperCAmelCase         =			ViltForImageAndTextRetrieval(A  )
						elif "mlm_itm" in checkpoint_url:
												UpperCAmelCase         =			True
												UpperCAmelCase         =			ViltForMaskedLM(A  )
						else:
												raise ValueError('''Unknown model type'''  )

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

						# Define processor
						UpperCAmelCase         =			ViltImageProcessor(size=3_84  )
						UpperCAmelCase         =			BertTokenizer.from_pretrained('''bert-base-uncased'''  )
						UpperCAmelCase         =			ViltProcessor(A   ,     A  )

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

						# Verify outputs
						if mlm_model:
												UpperCAmelCase         =			torch.Size([1, 11, 3_05_22]  )
												UpperCAmelCase         =			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]   ,     A   ,     atol=1E-4  )

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

	_lowercase    :    List[str]			     =  parser.parse_args()
	convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









'''simple docstring'''



# coding=utf-8
# Copyright 2020 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# this script dumps information about the environment

import os
import sys

import transformers


_lowercase    :    Union[str, Any]			     =  """3"""

print("""Python version:""", sys.version)
print("""transformers version:""", transformers.__version__)

try:
	import torch

	print("""Torch version:""", torch.__version__)
	print("""Cuda available:""", torch.cuda.is_available())
	print("""Cuda version:""", torch.version.cuda)
	print("""CuDNN version:""", torch.backends.cudnn.version())
	print("""Number of GPUs available:""", torch.cuda.device_count())
	print("""NCCL version:""", torch.cuda.nccl.version())
except ImportError:
	print("""Torch version:""", None)

try:
	import deepspeed

	print("""DeepSpeed version:""", deepspeed.__version__)
except ImportError:
	print("""DeepSpeed version:""", None)

try:
	import tensorflow as tf

	print("""TensorFlow version:""", tf.__version__)
	print("""TF GPUs available:""", bool(tf.config.list_physical_devices("""GPU""")))
	print("""Number of TF GPUs available:""", len(tf.config.list_physical_devices("""GPU""")))
except ImportError:
	print("""TensorFlow version:""", None)





