






'''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						:						List[Any]  =					logging.get_logger(__name__)
def __UpperCAmelCase   (					A					:			List[str]			,       A					:			List[str]=False			,       A					:			Tuple=False			,       A					:			List[Any]=False		)    ->   Tuple:
			UpperCAmelCase_      :       Optional[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   (					A					:			Tuple			,       A					:			Any		)    ->   List[Any]:
			for i in range(config.num_hidden_layers		):
						UpperCAmelCase_      :       Optional[Any]          =       '''vilt.'''
						# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
						UpperCAmelCase_      :       int          =       state_dict.pop(F"transformer.blocks.{i}.attn.qkv.weight"		)
						UpperCAmelCase_      :       Optional[int]          =       state_dict.pop(F"transformer.blocks.{i}.attn.qkv.bias"		)
						# next, add query, keys and values (in that order) to the state dict
						UpperCAmelCase_      :       Tuple          =       in_proj_weight[
						    : config.hidden_size, :
						]
						UpperCAmelCase_      :       int          =       in_proj_bias[: config.hidden_size]
						UpperCAmelCase_      :       Optional[int]          =       in_proj_weight[
						    config.hidden_size : config.hidden_size * 2, :
						]
						UpperCAmelCase_      :       Optional[int]          =       in_proj_bias[
						    config.hidden_size : config.hidden_size * 2
						]
						UpperCAmelCase_      :       Tuple          =       in_proj_weight[
						    -config.hidden_size :, :
						]
						UpperCAmelCase_      :       Optional[Any]          =       in_proj_bias[-config.hidden_size :]
def __UpperCAmelCase   (					A					:			int		)    ->   List[Any]:
			UpperCAmelCase_      :       Tuple          =       ['''head.weight''', '''head.bias''']
			for k in ignore_keys:
						state_dict.pop(A			,       A		)
def __UpperCAmelCase   (					A					:			Tuple			,       A					:			int			,       A					:			Any		)    ->   Tuple:
			UpperCAmelCase_      :       Any          =       dct.pop(A		)
			UpperCAmelCase_      :       Dict          =       val
@torch.no_grad()
def __UpperCAmelCase   (					A					:			Dict			,       A					:			Any		)    ->   Tuple:
			UpperCAmelCase_      :       Dict          =       ViltConfig(image_size=3_8_4			,       patch_size=3_2			,       tie_word_embeddings=A		)
			UpperCAmelCase_      :       int          =       False
			UpperCAmelCase_      :       Union[str, Any]          =       False
			UpperCAmelCase_      :       Optional[int]          =       False
			UpperCAmelCase_      :       Tuple          =       False
			if "vqa" in checkpoint_url:
						UpperCAmelCase_      :       Union[str, Any]          =       True
						UpperCAmelCase_      :       Dict          =       3_1_2_9
						UpperCAmelCase_      :       Union[str, Any]          =       '''huggingface/label-files'''
						UpperCAmelCase_      :       Optional[Any]          =       '''vqa2-id2label.json'''
						UpperCAmelCase_      :       List[Any]          =       json.load(open(hf_hub_download(A			,       A			,       repo_type='''dataset'''		)			,       '''r'''		)		)
						UpperCAmelCase_      :       str          =       {int(A		): v for k, v in idalabel.items()}
						UpperCAmelCase_      :       int          =       idalabel
						UpperCAmelCase_      :       Optional[int]          =       {v: k for k, v in idalabel.items()}
						UpperCAmelCase_      :       Optional[int]          =       ViltForQuestionAnswering(A		)
			elif "nlvr" in checkpoint_url:
						UpperCAmelCase_      :       Dict          =       True
						UpperCAmelCase_      :       Optional[Any]          =       2
						UpperCAmelCase_      :       List[str]          =       {0: '''False''', 1: '''True'''}
						UpperCAmelCase_      :       Optional[int]          =       {v: k for k, v in config.idalabel.items()}
						UpperCAmelCase_      :       Optional[int]          =       3
						UpperCAmelCase_      :       str          =       ViltForImagesAndTextClassification(A		)
			elif "irtr" in checkpoint_url:
						UpperCAmelCase_      :       Optional[Any]          =       True
						UpperCAmelCase_      :       List[Any]          =       ViltForImageAndTextRetrieval(A		)
			elif "mlm_itm" in checkpoint_url:
						UpperCAmelCase_      :       int          =       True
						UpperCAmelCase_      :       int          =       ViltForMaskedLM(A		)
			else:
						raise ValueError('''Unknown model type'''		)

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

			# Define processor
			UpperCAmelCase_      :       Optional[Any]          =       ViltImageProcessor(size=3_8_4		)
			UpperCAmelCase_      :       int          =       BertTokenizer.from_pretrained('''bert-base-uncased'''		)
			UpperCAmelCase_      :       Optional[int]          =       ViltProcessor(A			,       A		)

			# Forward pass on example inputs (image + text)
			if nlvr_model:
						UpperCAmelCase_      :       Union[str, Any]          =       Image.open(requests.get('''https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'''			,       stream=A		).raw		)
						UpperCAmelCase_      :       Optional[Any]          =       Image.open(requests.get('''https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'''			,       stream=A		).raw		)
						UpperCAmelCase_      :       List[str]          =       (
						    '''The left image contains twice the number of dogs as the right image, and at least two dogs in total are'''
						    ''' standing.'''
						)
						UpperCAmelCase_      :       Union[str, Any]          =       processor(A			,       A			,       return_tensors='''pt'''		)
						UpperCAmelCase_      :       Tuple          =       processor(A			,       A			,       return_tensors='''pt'''		)
						UpperCAmelCase_      :       List[str]          =       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=A		).raw		)
						if mlm_model:
									UpperCAmelCase_      :       int          =       '''a bunch of [MASK] laying on a [MASK].'''
						else:
									UpperCAmelCase_      :       Any          =       '''How many cats are there?'''
						UpperCAmelCase_      :       Tuple          =       processor(A			,       A			,       return_tensors='''pt'''		)
						UpperCAmelCase_      :       List[str]          =       model(**A		)

			# Verify outputs
			if mlm_model:
						UpperCAmelCase_      :       List[Any]          =       torch.Size([1, 1_1, 3_0_5_2_2]		)
						UpperCAmelCase_      :       List[str]          =       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_      :       str          =       outputs.logits[0, 4, :].argmax(-1		).item()
						assert tokenizer.decode([predicted_id]		) == "cats"
			elif vqa_model:
						UpperCAmelCase_      :       Tuple          =       torch.Size([1, 3_1_2_9]		)
						UpperCAmelCase_      :       Tuple          =       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_      :       str          =       outputs.logits.argmax(-1		).item()
						assert model.config.idalabel[predicted_idx] == "2"
			elif nlvr_model:
						UpperCAmelCase_      :       int          =       torch.Size([1, 2]		)
						UpperCAmelCase_      :       List[str]          =       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						:						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.'
						)

						_UpperCamelCase						:						Optional[Any]  =					parser.parse_args()
						convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)










'''simple docstring'''





from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


_UpperCamelCase						:						int  =					logging.get_logger(__name__)
def __UpperCAmelCase   (					A					:			Union[tf.Tensor, np.ndarray]		)    ->   List[int]:
			if isinstance(A			,       np.ndarray		):
						return list(tensor.shape		)

			UpperCAmelCase_      :       str          =       tf.shape(A		)

			if tensor.shape == tf.TensorShape(A		):
						return dynamic

			UpperCAmelCase_      :       Any          =       tensor.shape.as_list()

			return [dynamic[i] if s is None else s for i, s in enumerate(A		)]
def __UpperCAmelCase   (					A					:			tf.Tensor			,       A					:			Optional[int] = None			,       A					:			Optional[str] = None		)    ->   tf.Tensor:
			return tf.nn.softmax(logits=logits + 1e-9			,       axis=A			,       name=A		)
def __UpperCAmelCase   (					A					:			List[str]			,       A					:			Union[str, Any]			,       A					:			List[str]			,       A					:			str=1e-5			,       A					:			Dict=-1		)    ->   List[Any]:
			# This is a very simplified functional layernorm, designed to duplicate
			# the functionality of PyTorch nn.functional.layer_norm when this is needed to port
			# models in Transformers.

			if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(A			,       A		):
						raise NotImplementedError('''Only 1D weight and bias tensors are supported for now, with only a single axis.'''		)

			# Get mean and variance on the axis to be normalized
			UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       tf.nn.moments(A			,       axes=[axis]			,       keepdims=A		)

			if axis != -1:
						# Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
						# on every dimension except axis
						UpperCAmelCase_      :       Tuple          =       [1] * inputs.shape.rank
						UpperCAmelCase_      :       int          =       shape_list(A		)[axis]
						UpperCAmelCase_      :       Dict          =       tf.reshape(A			,       A		)
						UpperCAmelCase_      :       Optional[Any]          =       tf.reshape(A			,       A		)

			# Compute layer normalization using the batch_normalization
			# function.
			UpperCAmelCase_      :       Dict          =       tf.nn.batch_normalization(
			    A			,       A			,       A			,       offset=A			,       scale=A			,       variance_epsilon=A			,       )
			return outputs
def __UpperCAmelCase   (					A					:			Dict			,       A					:			int=0			,       A					:			Optional[int]=-1		)    ->   str:
			# Replicates the behavior of torch.flatten in TF

			# If end_dim or start_dim is negative, count them from the end
			if end_dim < 0:
						end_dim += input.shape.rank
			if start_dim < 0:
						start_dim += input.shape.rank

			if start_dim == end_dim:
						return input

			UpperCAmelCase_      :       Any          =       tf.shape(A		)
			UpperCAmelCase_      :       Dict          =       tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]		)
			UpperCAmelCase_      :       Optional[int]          =       tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]			,       axis=0		)
			return tf.reshape(A			,       A		)
def __UpperCAmelCase   (					A					:			tf.Tensor		)    ->   tf.Tensor:
			if not isinstance(A			,       tf.Tensor		):
						UpperCAmelCase_      :       Optional[Any]          =       tf.convert_to_tensor(A		)  # Catches stray NumPy inputs
			if encoder_attention_mask.shape.rank == 3:
						UpperCAmelCase_      :       Any          =       encoder_attention_mask[:, None, :, :]
			if encoder_attention_mask.shape.rank == 2:
						UpperCAmelCase_      :       Tuple          =       encoder_attention_mask[:, None, None, :]
			# T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
			# Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
			# /transformer/transformer_layers.py#L270
			# encoder_extended_attention_mask = (encoder_extended_attention_mask ==
			# encoder_extended_attention_mask.transpose(-1, -2))
			UpperCAmelCase_      :       int          =       (
			    tf.cast(1			,       encoder_attention_mask.dtype		) - encoder_extended_attention_mask
			) * encoder_extended_attention_mask.dtype.min

			return encoder_extended_attention_mask
def __UpperCAmelCase   (					A					:			tf.Tensor			,       A					:			int			,       A					:			str = "input_ids"		)    ->   None:
			tf.debugging.assert_less(
			    A			,       tf.cast(A			,       dtype=tensor.dtype		)			,       message=(
			        F"The maximum value of {tensor_name} ({tf.math.reduce_max(A		)}) must be smaller than the embedding "
			        F"layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."
			    )			,       )
def __UpperCAmelCase   (					A					:			Dict			,       A					:			Tuple			,       A					:			List[str]		)    ->   Union[str, Any]:
			UpperCAmelCase_      :       List[Any]          =       6_4_5_1_2
			# Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
			# because in that case even chunking the array would not make the saving
			# possible.
			UpperCAmelCase_      :       Tuple          =       [x for x in data if len(A		) > HDF5_OBJECT_HEADER_LIMIT]

			# Expecting this to never be true.
			if bad_attributes:
						raise RuntimeError(
						    '''The following attributes cannot be saved to HDF5 file because '''
						    F"they are larger than {HDF5_OBJECT_HEADER_LIMIT} "
						    F"bytes: {bad_attributes}"		)

			UpperCAmelCase_      :       Union[str, Any]          =       np.asarray(A		)

			UpperCAmelCase_      :       Optional[int]          =       1
			UpperCAmelCase_      :       str          =       np.array_split(A			,       A		)

			# This will never loop forever thanks to the test above.
			while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data		):
						num_chunks += 1
						UpperCAmelCase_      :       Dict          =       np.array_split(A			,       A		)

			if num_chunks > 1:
						for chunk_id, chunk_data in enumerate(A		):
									UpperCAmelCase_      :       int          =       chunk_data
			else:
						UpperCAmelCase_      :       List[Any]          =       data
def __UpperCAmelCase   (					A					:			int			,       A					:			Optional[Any]		)    ->   Tuple:
			if name in group.attrs:
						UpperCAmelCase_      :       Optional[int]          =       [n.decode('''utf8'''		) if hasattr(A			,       '''decode'''		) else n for n in group.attrs[name]]
			else:
						UpperCAmelCase_      :       Optional[int]          =       []
						UpperCAmelCase_      :       Optional[int]          =       0
						while "%s%d" % (name, chunk_id) in group.attrs:
									data.extend(
									    [n.decode('''utf8'''		) if hasattr(A			,       '''decode'''		) else n for n in group.attrs['''%s%d''' % (name, chunk_id)]]		)
									chunk_id += 1
			return data







def __UpperCAmelCase   (					A					:			Tuple		)    ->   str:

			def _expand_single_ad_tensor(A					:			int		):
						if isinstance(A			,       tf.Tensor		) and t.shape.rank == 1:
									return tf.expand_dims(A			,       axis=-1		)
						return t

			return tf.nest.map_structure(_expand_single_ad_tensor			,       A		)


