"""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()
UpperCAmelCase       =  logging.get_logger(__name__)






def      lowercase		(						a__	:		Dict ,  a__	:		int=False ,  a__	:		int=False ,  a__	:		int=False     )						->   Tuple:
					_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      lowercase		(						a__	:		List[str] ,  a__	:		List[Any]     )						->   List[str]:
					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      lowercase		(						a__	:		int     )						->   Any:
					_UpperCamelCase      						=	['head.weight', 'head.bias']
					for k in ignore_keys:
										state_dict.pop(a__ ,  a__     )






def      lowercase		(						a__	:		int ,  a__	:		List[Any] ,  a__	:		List[str]     )						->   Dict:
					_UpperCamelCase      						=	dct.pop(a__     )
					_UpperCamelCase      						=	val






@torch.no_grad()
def      lowercase		(						a__	:		int ,  a__	:		Union[str, Any]     )						->   Any:
					_UpperCamelCase      						=	ViltConfig(image_size=384 ,  patch_size=32 ,  tie_word_embeddings=a__     )
					_UpperCamelCase      						=	False
					_UpperCamelCase      						=	False
					_UpperCamelCase      						=	False
					_UpperCamelCase      						=	False
					if "vqa" in checkpoint_url:
										_UpperCamelCase      						=	True
										_UpperCamelCase      						=	3129
										_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      						=	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=384     )
					_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, 30522]     )
										_UpperCamelCase      						=	torch.tensor([-12.5061, -12.5123, -12.5174]     )
										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, 3129]     )
										_UpperCamelCase      						=	torch.tensor([-15.9495, -18.1472, -10.3041]     )
										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.8721, 2.1291]     )
										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__":
		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)





import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
 import torch

if is_vision_available():
 from PIL import Image

 from transformers import PoolFormerImageProcessor




class 					snake_case__						(       unittest.TestCase     ):






     """simple docstring"""


     def __init__(       self  :  Optional[int], _snake_case  :  List[Any], _snake_case  :  str=7, _snake_case  :  Tuple=3, _snake_case  :  List[str]=3_0, _snake_case  :  Tuple=4_0_0, _snake_case  :  Any=True, _snake_case  :  List[Any]=None, _snake_case  :  int=0.9, _snake_case  :  Optional[Any]=None, _snake_case  :  str=True, _snake_case  :  Union[str, Any]=[0.5, 0.5, 0.5], _snake_case  :  Union[str, Any]=[0.5, 0.5, 0.5], )   ->List[Any]:
            snake_case__							:	int   	=	size if size is not None else {'shortest_edge': 3_0}
            snake_case__							:	Tuple   	=	crop_size if crop_size is not None else {'height': 3_0, 'width': 3_0}
            snake_case__							:	Union[str, Any]   	=	parent
            snake_case__							:	Dict   	=	batch_size
            snake_case__							:	int   	=	num_channels
            snake_case__							:	Tuple   	=	min_resolution
            snake_case__							:	Any   	=	max_resolution
            snake_case__							:	List[Any]   	=	do_resize_and_center_crop
            snake_case__							:	str   	=	size
            snake_case__							:	str   	=	crop_pct
            snake_case__							:	List[str]   	=	crop_size
            snake_case__							:	Optional[int]   	=	do_normalize
            snake_case__							:	Tuple   	=	image_mean
            snake_case__							:	Tuple   	=	image_std






     def    lowercase_   (       self  :  Optional[int]					)   ->int:
            return {
                "size": self.size,
                "do_resize_and_center_crop": self.do_resize_and_center_crop,
                "crop_pct": self.crop_pct,
                "crop_size": self.crop_size,
                "do_normalize": self.do_normalize,
                "image_mean": self.image_mean,
                "image_std": self.image_std,
            }




@require_torch
@require_vision
class 					snake_case__						(       lowerCAmelCase_				,     unittest.TestCase     ):






     """simple docstring"""


     _SCREAMING_SNAKE_CASE						=	PoolFormerImageProcessor if is_vision_available() else None






     def    lowercase_   (       self  :  Union[str, Any]					)   ->Dict:
            snake_case__							:	Union[str, Any]   	=	PoolFormerImageProcessingTester(self					)






     @property
     def    lowercase_   (       self  :  int					)   ->Dict:
            return self.image_processor_tester.prepare_image_processor_dict()






     def    lowercase_   (       self  :  Union[str, Any]					)   ->Optional[int]:
            snake_case__							:	List[str]   	=	self.image_processing_class(**self.image_processor_dict					)
            self.assertTrue(hasattr(_snake_case, 'do_resize_and_center_crop'					)					)
            self.assertTrue(hasattr(_snake_case, 'size'					)					)
            self.assertTrue(hasattr(_snake_case, 'crop_pct'					)					)
            self.assertTrue(hasattr(_snake_case, 'do_normalize'					)					)
            self.assertTrue(hasattr(_snake_case, 'image_mean'					)					)
            self.assertTrue(hasattr(_snake_case, 'image_std'					)					)






     def    lowercase_   (       self  :  List[str]					)   ->List[str]:
            snake_case__							:	Union[str, Any]   	=	self.image_processing_class.from_dict(self.image_processor_dict					)
            self.assertEqual(image_processor.size, {'shortest_edge': 3_0}					)
            self.assertEqual(image_processor.crop_size, {'height': 3_0, 'width': 3_0}					)

            snake_case__							:	int   	=	self.image_processing_class.from_dict(self.image_processor_dict, size=4_2, crop_size=8_4					)
            self.assertEqual(image_processor.size, {'shortest_edge': 4_2}					)
            self.assertEqual(image_processor.crop_size, {'height': 8_4, 'width': 8_4}					)






     def    lowercase_   (       self  :  List[Any]					)   ->List[Any]:
            pass






     def    lowercase_   (       self  :  List[str]					)   ->str:
            # Initialize image_processing
            snake_case__							:	Union[str, Any]   	=	self.image_processing_class(**self.image_processor_dict					)
            # create random PIL images
            snake_case__							:	List[str]   	=	prepare_image_inputs(self.image_processor_tester, equal_resolution=_snake_case					)
            for image in image_inputs:
                   self.assertIsInstance(_snake_case, Image.Image					)

            # Test not batched input
            snake_case__							:	Optional[int]   	=	image_processing(image_inputs[0], return_tensors='pt'					).pixel_values

            self.assertEqual(
                encoded_images.shape, (
                    1,
                    self.image_processor_tester.num_channels,
                    self.image_processor_tester.crop_size['height'],
                    self.image_processor_tester.crop_size['width'],
                ), )

            # Test batched
            snake_case__							:	str   	=	image_processing(_snake_case, return_tensors='pt'					).pixel_values
            self.assertEqual(
                encoded_images.shape, (
                    self.image_processor_tester.batch_size,
                    self.image_processor_tester.num_channels,
                    self.image_processor_tester.crop_size['height'],
                    self.image_processor_tester.crop_size['width'],
                ), )






     def    lowercase_   (       self  :  int					)   ->List[Any]:
            # Initialize image_processing
            snake_case__							:	Optional[int]   	=	self.image_processing_class(**self.image_processor_dict					)
            # create random numpy tensors
            snake_case__							:	Dict   	=	prepare_image_inputs(self.image_processor_tester, equal_resolution=_snake_case, numpify=_snake_case					)
            for image in image_inputs:
                   self.assertIsInstance(_snake_case, np.ndarray					)

            # Test not batched input
            snake_case__							:	Dict   	=	image_processing(image_inputs[0], return_tensors='pt'					).pixel_values
            self.assertEqual(
                encoded_images.shape, (
                    1,
                    self.image_processor_tester.num_channels,
                    self.image_processor_tester.crop_size['height'],
                    self.image_processor_tester.crop_size['width'],
                ), )

            # Test batched
            snake_case__							:	List[Any]   	=	image_processing(_snake_case, return_tensors='pt'					).pixel_values
            self.assertEqual(
                encoded_images.shape, (
                    self.image_processor_tester.batch_size,
                    self.image_processor_tester.num_channels,
                    self.image_processor_tester.crop_size['height'],
                    self.image_processor_tester.crop_size['width'],
                ), )






     def    lowercase_   (       self  :  List[str]					)   ->List[str]:
            # Initialize image_processing
            snake_case__							:	Tuple   	=	self.image_processing_class(**self.image_processor_dict					)
            # create random PyTorch tensors
            snake_case__							:	List[str]   	=	prepare_image_inputs(self.image_processor_tester, equal_resolution=_snake_case, torchify=_snake_case					)
            for image in image_inputs:
                   self.assertIsInstance(_snake_case, torch.Tensor					)

            # Test not batched input
            snake_case__							:	Tuple   	=	image_processing(image_inputs[0], return_tensors='pt'					).pixel_values
            self.assertEqual(
                encoded_images.shape, (
                    1,
                    self.image_processor_tester.num_channels,
                    self.image_processor_tester.crop_size['height'],
                    self.image_processor_tester.crop_size['width'],
                ), )

            # Test batched
            snake_case__							:	Optional[Any]   	=	image_processing(_snake_case, return_tensors='pt'					).pixel_values
            self.assertEqual(
                encoded_images.shape, (
                    self.image_processor_tester.batch_size,
                    self.image_processor_tester.num_channels,
                    self.image_processor_tester.crop_size['height'],
                    self.image_processor_tester.crop_size['width'],
                ), )






