

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_									=   logging.get_logger(__name__)
def      a      (			A__	:   Optional[Any]      ,   A__	:   Tuple=False      ,   A__	:   Optional[int]=False      ,   A__	:   List[Any]=False     )    ->      List[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      a      (			A__	:   int      ,   A__	:   Dict     )    ->      int:





      """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      a      (			A__	:   Tuple     )    ->      Optional[Any]:





      """simple docstring"""



      _lowercase	       =['head.weight', 'head.bias']
      for k in ignore_keys:
            state_dict.pop(A__      ,   A__     )
def      a      (			A__	:   str      ,   A__	:   Optional[int]      ,   A__	:   List[Any]     )    ->      Optional[int]:





      """simple docstring"""



      _lowercase	       =dct.pop(A__     )
      _lowercase	       =val
@torch.no_grad()
def      a      (			A__	:   Union[str, Any]      ,   A__	:   Union[str, Any]     )    ->      str:





      """simple docstring"""



      _lowercase	       =ViltConfig(image_size=384      ,   patch_size=32      ,   tie_word_embeddings=A__     )
      _lowercase	       =False
      _lowercase	       =False
      _lowercase	       =False
      _lowercase	       =False
      if "vqa" in checkpoint_url:
            _lowercase	       =True
            _lowercase	       =3129
            _lowercase	       ='huggingface/label-files'
            _lowercase	       ='vqa2-id2label.json'
            _lowercase	       =json.load(open(hf_hub_download(A__      ,   A__      ,   repo_type='dataset'     )      ,   'r'     )     )
            _lowercase	       ={int(A__     ): v for k, v in idalabel.items()}
            _lowercase	       =idalabel
            _lowercase	       ={v: k for k, v in idalabel.items()}
            _lowercase	       =ViltForQuestionAnswering(A__     )
      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(A__     )
      elif "irtr" in checkpoint_url:
            _lowercase	       =True
            _lowercase	       =ViltForImageAndTextRetrieval(A__     )
      elif "mlm_itm" in checkpoint_url:
            _lowercase	       =True
            _lowercase	       =ViltForMaskedLM(A__     )
      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(A__      ,   map_location='cpu'     )['state_dict']
      _lowercase	       =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:
            _lowercase	       =['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:
            _lowercase       ,       _lowercase	       =model.load_state_dict(A__      ,   strict=A__     )
            assert missing_keys == ["mlm_score.decoder.bias"]
      else:
            model.load_state_dict(A__     )

      # Define processor
      _lowercase	       =ViltImageProcessor(size=384     )
      _lowercase	       =BertTokenizer.from_pretrained('bert-base-uncased'     )
      _lowercase	       =ViltProcessor(A__      ,   A__     )

      # 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=A__     ).raw     )
            _lowercase	       =Image.open(requests.get('https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'      ,   stream=A__     ).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(A__      ,   A__      ,   return_tensors='pt'     )
            _lowercase	       =processor(A__      ,   A__      ,   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=A__     ).raw     )
            if mlm_model:
                  _lowercase	       ='a bunch of [MASK] laying on a [MASK].'
            else:
                  _lowercase	       ='How many cats are there?'
            _lowercase	       =processor(A__      ,   A__      ,   return_tensors='pt'     )
            _lowercase	       =model(**A__     )

      # Verify outputs
      if mlm_model:
            _lowercase	       =torch.Size([1, 11, 30522]     )
            _lowercase	       =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"
            _lowercase	       =outputs.logits[0, 4, :].argmax(-1     ).item()
            assert tokenizer.decode([predicted_id]     ) == "cats"
      elif vqa_model:
            _lowercase	       =torch.Size([1, 3129]     )
            _lowercase	       =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"
            _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]      ,   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_									=   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_									=   parser.parse_args()
  convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


lowercase_									=   logging.getLogger(__name__)





class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase=None							,   lowerCAmelCase=None					)  -> List[str]:
        '''simple docstring'''






        _lowercase	       =self.layer[current_layer](lowerCAmelCase							,   lowerCAmelCase							,   head_mask[current_layer]					)

        _lowercase	       =layer_outputs[0]

        return hidden_states





@add_start_docstrings(
    """The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top."""			,			SCREAMING_SNAKE_CASE			,			)
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def __init__(				self							,   lowerCAmelCase					)  -> Union[str, Any]:
        '''simple docstring'''






        super().__init__(lowerCAmelCase					)

        _lowercase	       =BertEncoderWithPabee(lowerCAmelCase					)

        self.init_weights()
        _lowercase	       =0
        _lowercase	       =0
        _lowercase	       =0

        _lowercase	       =0




  def   A__     (				self							,   lowerCAmelCase					)  -> Optional[Any]:
        '''simple docstring'''






        _lowercase	       =threshold




  def   A__     (				self							,   lowerCAmelCase					)  -> List[Any]:
        '''simple docstring'''






        _lowercase	       =patience




  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        _lowercase	       =0
        _lowercase	       =0




  def   A__     (				self					)  -> int:
        '''simple docstring'''






        _lowercase	       =self.inference_layers_num / self.inference_instances_num
        _lowercase	       =(
            F'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
            F''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
        )
        print(lowerCAmelCase					)




  @add_start_docstrings_to_model_forward(lowerCAmelCase					)
  def   A__     (				self							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=False							,   )  -> str:
        '''simple docstring'''






        if input_ids is not None and inputs_embeds is not None:
              raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time'					)
        elif input_ids is not None:
              _lowercase	       =input_ids.size()
        elif inputs_embeds is not None:
              _lowercase	       =inputs_embeds.size()[:-1]
        else:
              raise ValueError('You have to specify either input_ids or inputs_embeds'					)

        _lowercase	       =input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
              _lowercase	       =torch.ones(lowerCAmelCase							,   device=lowerCAmelCase					)
        if token_type_ids is None:
              _lowercase	       =torch.zeros(lowerCAmelCase							,   dtype=torch.long							,   device=lowerCAmelCase					)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        _lowercase	       =self.get_extended_attention_mask(lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					)

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
              _lowercase       ,       _lowercase       ,       _lowercase	       =encoder_hidden_states.size()
              _lowercase	       =(encoder_batch_size, encoder_sequence_length)
              if encoder_attention_mask is None:
                    _lowercase	       =torch.ones(lowerCAmelCase							,   device=lowerCAmelCase					)
              _lowercase	       =self.invert_attention_mask(lowerCAmelCase					)
        else:
              _lowercase	       =None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        _lowercase	       =self.get_head_mask(lowerCAmelCase							,   self.config.num_hidden_layers					)

        _lowercase	       =self.embeddings(
            input_ids=lowerCAmelCase							,   position_ids=lowerCAmelCase							,   token_type_ids=lowerCAmelCase							,   inputs_embeds=lowerCAmelCase					)
        _lowercase	       =embedding_output

        if self.training:
              _lowercase	       =[]
              for i in range(self.config.num_hidden_layers					):
                    _lowercase	       =self.encoder.adaptive_forward(
                        lowerCAmelCase							,   current_layer=lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   head_mask=lowerCAmelCase					)

                    _lowercase	       =self.pooler(lowerCAmelCase					)
                    _lowercase	       =output_layers[i](output_dropout(lowerCAmelCase					)					)
                    res.append(lowerCAmelCase					)
        elif self.patience == 0:  # Use all layers for inference
              _lowercase	       =self.encoder(
                  lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   head_mask=lowerCAmelCase							,   encoder_hidden_states=lowerCAmelCase							,   encoder_attention_mask=lowerCAmelCase							,   )
              _lowercase	       =self.pooler(encoder_outputs[0]					)
              _lowercase	       =[output_layers[self.config.num_hidden_layers - 1](lowerCAmelCase					)]
        else:
              _lowercase	       =0
              _lowercase	       =None
              _lowercase	       =0
              for i in range(self.config.num_hidden_layers					):
                    calculated_layer_num += 1
                    _lowercase	       =self.encoder.adaptive_forward(
                        lowerCAmelCase							,   current_layer=lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   head_mask=lowerCAmelCase					)

                    _lowercase	       =self.pooler(lowerCAmelCase					)
                    _lowercase	       =output_layers[i](lowerCAmelCase					)
                    if regression:
                          _lowercase	       =logits.detach()
                          if patient_result is not None:
                                _lowercase	       =patient_result.detach()
                          if (patient_result is not None) and torch.abs(patient_result - labels					) < self.regression_threshold:
                                patient_counter += 1
                          else:
                                _lowercase	       =0
                    else:
                          _lowercase	       =logits.detach().argmax(dim=1					)
                          if patient_result is not None:
                                _lowercase	       =patient_result.detach().argmax(dim=1					)
                          if (patient_result is not None) and torch.all(labels.eq(lowerCAmelCase					)					):
                                patient_counter += 1
                          else:
                                _lowercase	       =0

                    _lowercase	       =logits
                    if patient_counter == self.patience:
                          break
              _lowercase	       =[patient_result]
              self.inference_layers_num += calculated_layer_num
              self.inference_instances_num += 1

        return res





@add_start_docstrings(
    """Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. """			,			SCREAMING_SNAKE_CASE			,			)
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def __init__(				self							,   lowerCAmelCase					)  -> List[Any]:
        '''simple docstring'''






        super().__init__(lowerCAmelCase					)
        _lowercase	       =config.num_labels

        _lowercase	       =BertModelWithPabee(lowerCAmelCase					)
        _lowercase	       =nn.Dropout(config.hidden_dropout_prob					)
        _lowercase	       =nn.ModuleList(
            [nn.Linear(config.hidden_size							,   self.config.num_labels					) for _ in range(config.num_hidden_layers					)]					)

        self.init_weights()




  @add_start_docstrings_to_model_forward(lowerCAmelCase					)
  def   A__     (				self							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   )  -> Union[str, Any]:
        '''simple docstring'''






        _lowercase	       =self.bert(
            input_ids=lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   token_type_ids=lowerCAmelCase							,   position_ids=lowerCAmelCase							,   head_mask=lowerCAmelCase							,   inputs_embeds=lowerCAmelCase							,   output_dropout=self.dropout							,   output_layers=self.classifiers							,   regression=self.num_labels == 1							,   )

        _lowercase	       =(logits[-1],)

        if labels is not None:
              _lowercase	       =None
              _lowercase	       =0
              for ix, logits_item in enumerate(lowerCAmelCase					):
                    if self.num_labels == 1:
                          #  We are doing regression
                          _lowercase	       =MSELoss()
                          _lowercase	       =loss_fct(logits_item.view(-1					)							,   labels.view(-1					)					)
                    else:
                          _lowercase	       =CrossEntropyLoss()
                          _lowercase	       =loss_fct(logits_item.view(-1							,   self.num_labels					)							,   labels.view(-1					)					)
                    if total_loss is None:
                          _lowercase	       =loss
                    else:
                          total_loss += loss * (ix + 1)
                    total_weights += ix + 1
              _lowercase	       =(total_loss / total_weights,) + outputs

        return outputs


