


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







def 			_lowerCamelCase(      lowercase__  ,					lowercase__=False  ,					lowercase__=False  ,					lowercase__=False      )		->      Union[str, Any]:

		'''simple docstring'''







		__lowercase=     []
		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(      lowercase__  ,					lowercase__      )		->      List[str]:

		'''simple docstring'''







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







def 			_lowerCamelCase(      lowercase__      )		->      str:

		'''simple docstring'''







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







def 			_lowerCamelCase(      lowercase__  ,					lowercase__  ,					lowercase__      )		->      Optional[Any]:

		'''simple docstring'''







		__lowercase=     dct.pop(lowercase__      )
		__lowercase=     val







@torch.no_grad()
def 			_lowerCamelCase(      lowercase__  ,					lowercase__      )		->      Tuple:

		'''simple docstring'''







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

		# load state_dict of original model, remove and rename some keys
		__lowercase=     torch.hub.load_state_dict_from_url(lowercase__  ,					map_location='cpu'      )['''state_dict''']
		__lowercase=     create_rename_keys(lowercase__  ,					lowercase__  ,					lowercase__  ,					lowercase__      )
		for src, dest in rename_keys:
				rename_key(lowercase__  ,					lowercase__  ,					lowercase__      )
		read_in_q_k_v(lowercase__  ,					lowercase__      )
		if mlm_model or irtr_model:
				__lowercase=     ['''itm_score.fc.weight''', '''itm_score.fc.bias''']
				for k in ignore_keys:
						state_dict.pop(lowercase__  ,					lowercase__      )

    # load state dict into HuggingFace model
		model.eval()
		if mlm_model:
				__lowercase=     model.load_state_dict(lowercase__  ,					strict=lowercase__      )
				assert missing_keys == ["mlm_score.decoder.bias"]
		else:
				model.load_state_dict(lowercase__      )

		# Define processor
		__lowercase=     ViltImageProcessor(size=3_8_4      )
		__lowercase=     BertTokenizer.from_pretrained('bert-base-uncased'      )
		__lowercase=     ViltProcessor(lowercase__  ,					lowercase__      )

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

		# Verify outputs
		if mlm_model:
				__lowercase=     torch.Size([1, 1_1, 3_0_5_2_2]      )
				__lowercase=     torch.tensor([-12.5061, -12.5123, -12.5174]      )
				assert outputs.logits.shape == expected_shape
				assert torch.allclose(outputs.logits[0, 0, :3]  ,					lowercase__  ,					atol=1E-4      )

				# verify masked token prediction equals "cats"
				__lowercase=     outputs.logits[0, 4, :].argmax(-1      ).item()
				assert tokenizer.decode([predicted_id]      ) == "cats"
		elif vqa_model:
				__lowercase=     torch.Size([1, 3_1_2_9]      )
				__lowercase=     torch.tensor([-15.9495, -18.1472, -10.3041]      )
				assert torch.allclose(outputs.logits[0, :3]  ,					lowercase__  ,					atol=1E-4      )
				assert outputs.logits.shape == expected_shape
				assert torch.allclose(outputs.logits[0, 0, :3]  ,					lowercase__  ,					atol=1E-4      )

				# verify vqa prediction equals "2"
				__lowercase=     outputs.logits.argmax(-1      ).item()
				assert model.config.idalabel[predicted_idx] == "2"
		elif nlvr_model:
				__lowercase=     torch.Size([1, 2]      )
				__lowercase=     torch.tensor([-2.8721, 2.1291]      )
				assert torch.allclose(outputs.logits[0, :3]  ,					lowercase__  ,					atol=1E-4      )
				assert outputs.logits.shape == expected_shape

		Path(lowercase__      ).mkdir(exist_ok=lowercase__      )
		print(F'Saving model and processor to {pytorch_dump_folder_path}'      )
		model.save_pretrained(lowercase__      )
		processor.save_pretrained(lowercase__      )


if __name__ == "__main__":
	lowerCAmelCase       			=						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.'''
	)

	lowerCAmelCase       			=						parser.parse_args()
	convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)














import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


SCREAMING_SNAKE_CASE__   :      List[str]				=		logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
					import torch_xla.core.xla_model as xm
					import torch_xla.debug.metrics as met
class     lowerCamelCase_					(   lowerCamelCase			):







				def __init__(    self							,       *__lowerCAmelCase							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       **__lowerCAmelCase  ):
								"""simple docstring"""
								super().__init__(*__lowerCAmelCase							,       **__lowerCAmelCase  )
								__magic_name__		:Any   =   eval_examples
								__magic_name__		:str   =   post_process_function
								__magic_name__		:int   =   quant_trainer_args
								__magic_name__		:List[str]   =   1_2_8  # default number of calibration samples







				def       A						(    self							,       __lowerCAmelCase=None  ):
								"""simple docstring"""
								if calib_dataset is None and self.calib_dataset is None:
												raise ValueError('''Trainer: calibration requires an calib_dataset.'''  )
								__magic_name__		:Optional[Any]   =   calib_dataset if calib_dataset is not None else self.calib_dataset

								__magic_name__		:Optional[int]   =   self._remove_unused_columns(__lowerCAmelCase							,       description='''Calibration'''  )

								return DataLoader(
								    __lowerCAmelCase							,       batch_size=self.args.eval_batch_size							,       collate_fn=self.data_collator							,       drop_last=self.args.dataloader_drop_last							,       num_workers=self.args.dataloader_num_workers							,       pin_memory=self.args.dataloader_pin_memory							,       shuffle=__lowerCAmelCase							,       )







				def       A						(    self							,       __lowerCAmelCase=None  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   self.train_dataset if calib_dataset is None else calib_dataset
								__magic_name__		:Any   =   self.get_calib_dataloader(__lowerCAmelCase  )

								__magic_name__		:List[str]   =   self.model
								quant_trainer.configure_model(__lowerCAmelCase							,       self.quant_trainer_args							,       calib=__lowerCAmelCase  )
								model.eval()
								quant_trainer.enable_calibration(__lowerCAmelCase  )

								logger.info('''***** Running calibration *****'''  )
								logger.info(F'''  Num examples = {self.calib_num}'''  )
								logger.info(F'''  Batch size = {calib_dataloader.batch_size}'''  )

								for step, inputs in enumerate(__lowerCAmelCase  ):
												# Prediction step
												__magic_name__  ,					__magic_name__  ,					__magic_name__		:str   =   self.prediction_step(__lowerCAmelCase							,       __lowerCAmelCase							,       prediction_loss_only=__lowerCAmelCase  )
												if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
																break

								quant_trainer.finish_calibration(__lowerCAmelCase							,       self.quant_trainer_args  )
								__magic_name__		:Any   =   model







				def       A						(    self							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       __lowerCAmelCase = "eval"  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   self.eval_dataset if eval_dataset is None else eval_dataset
								__magic_name__		:Optional[Any]   =   self.get_eval_dataloader(__lowerCAmelCase  )
								__magic_name__		:str   =   self.eval_examples if eval_examples is None else eval_examples

								# Temporarily disable metric computation, we will do it in the loop here.
								__magic_name__		:Any   =   self.compute_metrics
								__magic_name__		:List[Any]   =   None
								__magic_name__		:List[str]   =   self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								try:
												__magic_name__		:Optional[Any]   =   eval_loop(
												    __lowerCAmelCase							,       description='''Evaluation'''							,       prediction_loss_only=True if compute_metrics is None else None							,       ignore_keys=__lowerCAmelCase							,       )
								finally:
												__magic_name__		:Union[str, Any]   =   compute_metrics

								if self.post_process_function is not None and self.compute_metrics is not None:
												__magic_name__		:Union[str, Any]   =   self.post_process_function(__lowerCAmelCase							,       __lowerCAmelCase							,       output.predictions  )
												__magic_name__		:int   =   self.compute_metrics(__lowerCAmelCase  )

												# Prefix all keys with metric_key_prefix + '_'
												for key in list(metrics.keys()  ):
																if not key.startswith(F'''{metric_key_prefix}_'''  ):
																				__magic_name__		:Dict   =   metrics.pop(__lowerCAmelCase  )

												self.log(__lowerCAmelCase  )
								else:
												__magic_name__		:List[str]   =   {}

								if self.args.tpu_metrics_debug or self.args.debug:
												# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
												xm.master_print(met.metrics_report()  )

								__magic_name__		:Optional[Any]   =   self.callback_handler.on_evaluate(self.args							,       self.state							,       self.control							,       __lowerCAmelCase  )
								return metrics







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase=None							,       __lowerCAmelCase = "test"  ):
								"""simple docstring"""
								__magic_name__		:int   =   self.get_test_dataloader(__lowerCAmelCase  )

								# Temporarily disable metric computation, we will do it in the loop here.
								__magic_name__		:Dict   =   self.compute_metrics
								__magic_name__		:str   =   None
								__magic_name__		:Optional[int]   =   self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								try:
												__magic_name__		:int   =   eval_loop(
												    __lowerCAmelCase							,       description='''Prediction'''							,       prediction_loss_only=True if compute_metrics is None else None							,       ignore_keys=__lowerCAmelCase							,       )
								finally:
												__magic_name__		:List[Any]   =   compute_metrics

								if self.post_process_function is None or self.compute_metrics is None:
												return output

								__magic_name__		:Optional[Any]   =   self.post_process_function(__lowerCAmelCase							,       __lowerCAmelCase							,       output.predictions							,       '''predict'''  )
								__magic_name__		:Dict   =   self.compute_metrics(__lowerCAmelCase  )

								# Prefix all keys with metric_key_prefix + '_'
								for key in list(metrics.keys()  ):
												if not key.startswith(F'''{metric_key_prefix}_'''  ):
																__magic_name__		:List[str]   =   metrics.pop(__lowerCAmelCase  )

								return PredictionOutput(predictions=predictions.predictions							,       label_ids=predictions.label_ids							,       metrics=__lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase="./"  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   self.eval_dataset
								__magic_name__		:Any   =   self.get_eval_dataloader(__lowerCAmelCase  )

								__magic_name__		:int   =   next(iter(__lowerCAmelCase  )  )

								# saving device - to make it consistent
								__magic_name__		:str   =   torch.device('''cuda''' if torch.cuda.is_available() else '''cpu'''  )

								# convert to tuple
								__magic_name__		:int   =   tuple(v.to(__lowerCAmelCase  ) for k, v in batch.items()  )

								logger.info('''Converting model to be onnx compatible'''  )
								from pytorch_quantization.nn import TensorQuantizer

								__magic_name__		:Any   =   True

								__magic_name__		:Optional[int]   =   self.model.to(__lowerCAmelCase  )

								model.eval()
								model.float()

								__magic_name__		:Any   =   model.module if hasattr(__lowerCAmelCase							,       '''module'''  ) else model
								quant_trainer.configure_model(__lowerCAmelCase							,       self.quant_trainer_args  )

								__magic_name__		:int   =   os.path.join(__lowerCAmelCase							,       '''model.onnx'''  )
								logger.info(F'''exporting model to {output_model_file}'''  )

								__magic_name__		:Dict   =   {0: '''batch_size''', 1: '''seq_len'''}

								torch.onnx.export(
								    __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       export_params=__lowerCAmelCase							,       opset_version=1_3							,       do_constant_folding=__lowerCAmelCase							,       input_names=['''input_ids''', '''attention_mask''', '''token_type_ids''']							,       output_names=['''output_start_logits''', '''output_end_logits''']							,       dynamic_axes={
								        '''input_ids''': axes,
								        '''attention_mask''': axes,
								        '''token_type_ids''': axes,
								        '''output_start_logits''': axes,
								        '''output_end_logits''': axes,
								    }							,       verbose=__lowerCAmelCase							,       )
								logger.info('''onnx export finished'''  )






