




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   :						Optional[Any]				=			logging.get_logger(__name__)







def 		lowercase__   (		__A:	List[Any]							,__A:	List[Any]=False							,__A:	Tuple=False							,__A:	Dict=False				):





     '''simple docstring'''






     __magic_name__							:	Union[str, Any]    							=							[]
     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[Any]							,__A:	Any				):





     '''simple docstring'''






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







def 		lowercase__   (		__A:	Dict				):





     '''simple docstring'''






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







def 		lowercase__   (		__A:	List[Any]							,__A:	Tuple							,__A:	Any				):





     '''simple docstring'''






     __magic_name__							:	str    							=							dct.pop(__A				)
     __magic_name__							:	Optional[Any]    							=							val







@torch.no_grad()
def 		lowercase__   (		__A:	str							,__A:	Optional[int]				):





     '''simple docstring'''






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

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

     # Define processor
     __magic_name__							:	Tuple    							=							ViltImageProcessor(size=3_8_4				)
     __magic_name__							:	Union[str, Any]    							=							BertTokenizer.from_pretrained('''bert-base-uncased'''				)
     __magic_name__							:	Tuple    							=							ViltProcessor(__A							,__A				)

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

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

 __lowerCamelCase   :						List[Any]				=			parser.parse_args()
 convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCamelCase   :						Optional[int]				=			logging.get_logger(__name__)

__lowerCamelCase   :						Tuple				=			{
    '''shi-labs/nat-mini-in1k-224''': '''https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json''',
    # See all Nat models at https://huggingface.co/models?filter=nat
}


class    lowerCamelCase    (		_lowerCamelCase				,_lowerCamelCase						):






			'''simple docstring'''





			UpperCamelCase__								='''nat'''

			UpperCamelCase__								={
			    '''num_attention_heads''': '''num_heads''',
			    '''num_hidden_layers''': '''num_layers''',
			}

			def __init__( self			:   Union[str, Any]       ,						lowerCamelCase_			:   Dict=4       ,						lowerCamelCase_			:   Union[str, Any]=3       ,						lowerCamelCase_			:   List[str]=64       ,						lowerCamelCase_			:   Union[str, Any]=[3, 4, 6, 5]       ,						lowerCamelCase_			:   List[Any]=[2, 4, 8, 16]       ,						lowerCamelCase_			:   List[Any]=7       ,						lowerCamelCase_			:   Union[str, Any]=3.0       ,						lowerCamelCase_			:   int=True       ,						lowerCamelCase_			:   str=0.0       ,						lowerCamelCase_			:   int=0.0       ,						lowerCamelCase_			:   Dict=0.1       ,						lowerCamelCase_			:   Optional[int]="gelu"       ,						lowerCamelCase_			:   Tuple=0.0_2       ,						lowerCamelCase_			:   Tuple=1E-5       ,						lowerCamelCase_			:   List[Any]=0.0       ,						lowerCamelCase_			:   str=None       ,						lowerCamelCase_			:   Tuple=None       ,						**lowerCamelCase_			:   Any       ,						)						->   Any:
								super().__init__(**lowerCamelCase_  )

								__magic_name__							:	List[Any]    							=							patch_size
								__magic_name__							:	str    							=							num_channels
								__magic_name__							:	Union[str, Any]    							=							embed_dim
								__magic_name__							:	Dict    							=							depths
								__magic_name__							:	str    							=							len(lowerCamelCase_  )
								__magic_name__							:	List[str]    							=							num_heads
								__magic_name__							:	Optional[int]    							=							kernel_size
								__magic_name__							:	Any    							=							mlp_ratio
								__magic_name__							:	Any    							=							qkv_bias
								__magic_name__							:	Union[str, Any]    							=							hidden_dropout_prob
								__magic_name__							:	Any    							=							attention_probs_dropout_prob
								__magic_name__							:	List[str]    							=							drop_path_rate
								__magic_name__							:	str    							=							hidden_act
								__magic_name__							:	Tuple    							=							layer_norm_eps
								__magic_name__							:	Any    							=							initializer_range
								# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
								# this indicates the channel dimension after the last stage of the model
								__magic_name__							:	str    							=							int(embed_dim * 2 ** (len(lowerCamelCase_  ) - 1)  )
								__magic_name__							:	str    							=							layer_scale_init_value
								__magic_name__							:	Dict    							=							['''stem'''] + [F'''stage{idx}''' for idx in range(1       ,						len(lowerCamelCase_  ) + 1  )]
								__magic_name__							,    __magic_name__							:	int    							=							get_aligned_output_features_output_indices(
								    out_features=lowerCamelCase_       ,						out_indices=lowerCamelCase_       ,						stage_names=self.stage_names  )



