import logging
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import librosa
import torch
from datasets import DatasetDict, load_dataset
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForPreTraining,
    is_apex_available,
    trainer_utils,
)
from transformers.models.wavaveca.modeling_wavaveca import _compute_mask_indices


if is_apex_available():
      from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse("1.6"):
      SCREAMING_SNAKE_CASE__		:					int						=					True
      from torch.cuda.amp import autocast


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

@dataclass
class       snake_case					:
  lowercase_			=     field(
      metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}						)
  lowercase_			=     field(
      default=UpperCamelCase_   ,						metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}   ,						)
  lowercase_			=     field(
      default=UpperCamelCase_   ,						metadata={'help': 'Whether to freeze the feature extractor layers of the model.'}						)
  lowercase_			=     field(
      default=UpperCamelCase_   ,						metadata={'help': 'Whether to log verbose messages or not.'}   ,						)
  lowercase_			=     field(
      default=2.0   ,						metadata={'help': 'Maximum temperature for gumbel softmax.'}						)
  lowercase_			=     field(
      default=0.5   ,						metadata={'help': 'Minimum temperature for gumbel softmax.'}						)
  lowercase_			=     field(
      default=0.999995   ,						metadata={'help': 'Decay of gumbel temperature during training.'}						)







def     _a				(      lowercase__  :			ModelArguments     ,  lowercase__  :			TrainingArguments ):







     '''simple docstring'''


     logging.basicConfig(
         format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'     ,  datefmt='%m/%d/%Y %H:%M:%S'     ,  handlers=[logging.StreamHandler(sys.stdout )]     ,  )
     SCREAMING_SNAKE_CASE__				:				str			=		logging.WARNING
     if model_args.verbose_logging:
          SCREAMING_SNAKE_CASE__				:				Dict			=		logging.DEBUG
     elif trainer_utils.is_main_process(training_args.local_rank ):
          SCREAMING_SNAKE_CASE__				:				List[str]			=		logging.INFO
     logger.setLevel(lowercase__ )

@dataclass
class       snake_case					:
  lowercase_			=     field(
      default=UpperCamelCase_   ,						metadata={'help': 'The name of the dataset to use (via the datasets library).'}						)
  lowercase_			=     field(
      default=UpperCamelCase_   ,						metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'}						)
  lowercase_			=     field(
      default='train'   ,						metadata={
          'help': 'The name of the training data set split to use (via the datasets library). Defaults to \'train\''
      }   ,						)
  lowercase_			=     field(
      default='validation'   ,						metadata={
          'help': (
              'The name of the validation data set split to use (via the datasets library). Defaults to \'validation\''
          )
      }   ,						)
  lowercase_			=     field(
      default='file'   ,						metadata={'help': 'Column in the dataset that contains speech file path. Defaults to \'file\''}   ,						)
  lowercase_			=     field(
      default=UpperCamelCase_   ,						metadata={'help': 'Overwrite the cached preprocessed datasets or not.'}						)
  lowercase_			=     field(
      default=1   ,						metadata={
          'help': 'The percentage of the train set used as validation set in case there\'s no validation split'
      }   ,						)
  lowercase_			=     field(
      default=UpperCamelCase_   ,						metadata={'help': 'The number of processes to use for the preprocessing.'}   ,						)
  lowercase_			=     field(
      default=20.0   ,						metadata={'help': 'Filter audio files that are longer than `max_duration_in_seconds` seconds'}						)




@dataclass
class       snake_case					:
  lowercase_			=     42
  lowercase_			=     42
  lowercase_			=     "longest"
  lowercase_			=     None
  lowercase_			=     None




  def __call__(      self     :	int     , a_     :	List[Dict[str, Union[List[int], torch.Tensor]]]				)->       Dict[str, torch.Tensor]:




       """simple docstring"""
       # reformat list to dict and set to pytorch format
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		self.feature_extractor.pad(
           a_     , max_length=self.max_length     , padding=self.padding     , pad_to_multiple_of=self.pad_to_multiple_of     , return_tensors='pt'     , )
       SCREAMING_SNAKE_CASE__				:				List[Any]			=		self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]				)

       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		batch['input_values'].shape[0]

       # make sure that no loss is computed on padded inputs
       if batch["attention_mask"] is not None:
            # compute real output lengths according to convolution formula
            SCREAMING_SNAKE_CASE__				:				Tuple			=		self.model._get_feat_extract_output_lengths(batch['attention_mask'].sum(-1				)				).to(
                torch.long				)

            SCREAMING_SNAKE_CASE__				:				str			=		torch.zeros(
                (batch_size, mask_indices_seq_length)     , dtype=torch.long     , device=batch['input_values'].device				)

            # these two operations makes sure that all values
            # before the output lengths indices are attended to
            SCREAMING_SNAKE_CASE__				:				List[Any]			=		1
            SCREAMING_SNAKE_CASE__				:				Optional[int]			=		attention_mask.flip([-1]				).cumsum(-1				).flip([-1]				).bool()

       # sample randomly masked indices
       SCREAMING_SNAKE_CASE__				:				str			=		_compute_mask_indices(
           (batch_size, mask_indices_seq_length)     , self.model.config.mask_time_prob     , self.model.config.mask_time_length     , attention_mask=a_     , min_masks=2     , )

       return batch




class       snake_case					(      UpperCamelCase_						):




  def __init__(      self     :	str     , *a_     :	List[str]     , a_     :	Dict=1     , a_     :	Union[str, Any]=0     , a_     :	Union[str, Any]=1.0     , **a_     :	List[str]				)->       Any:




       """simple docstring"""
       super().__init__(*a_     , **a_				)
       SCREAMING_SNAKE_CASE__				:				Any			=		0
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		max_gumbel_temp
       SCREAMING_SNAKE_CASE__				:				int			=		min_gumbel_temp
       SCREAMING_SNAKE_CASE__				:				int			=		gumbel_temp_decay




  def 	__lowercase(      self     :	int     , a_     :	nn.Module     , a_     :	Dict[str, Union[torch.Tensor, Any]]				)->       torch.Tensor:




       """simple docstring"""
       model.train()
       SCREAMING_SNAKE_CASE__				:				Dict			=		self._prepare_inputs(a_				)

       if self.use_amp:
            with autocast():
                 SCREAMING_SNAKE_CASE__				:				List[str]			=		self.compute_loss(a_     , a_				)
       else:
            SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		self.compute_loss(a_     , a_				)

       if self.args.n_gpu > 1 or self.deepspeed:
            if model.module.config.ctc_loss_reduction == "mean":
                 SCREAMING_SNAKE_CASE__				:				str			=		loss.mean()
            elif model.module.config.ctc_loss_reduction == "sum":
                 SCREAMING_SNAKE_CASE__				:				List[Any]			=		loss.sum() / (inputs['mask_time_indices']).sum()
            else:
                 raise ValueError(F'''{model.config.ctc_loss_reduction} is not valid. Choose one of [\'mean\', \'sum\']'''				)

       if self.args.gradient_accumulation_steps > 1:
            SCREAMING_SNAKE_CASE__				:				str			=		loss / self.args.gradient_accumulation_steps

       if self.use_amp:
            self.scaler.scale(a_				).backward()
       elif self.use_apex:
            with amp.scale_loss(a_     , self.optimizer				) as scaled_loss:
                 scaled_loss.backward()
       elif self.deepspeed:
            self.deepspeed.backward(a_				)
       else:
            loss.backward()

       self.num_update_step += 1
       # make sure gumbel softmax temperature is decayed
       if self.args.n_gpu > 1 or self.deepspeed:
            model.module.set_gumbel_temperature(
                max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step     , self.min_gumbel_temp				)				)
       else:
            model.set_gumbel_temperature(
                max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step     , self.min_gumbel_temp				)				)

       return loss.detach()







def     _a				(      ):







     '''simple docstring'''


     SCREAMING_SNAKE_CASE__				:				int			=		HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments) )

     SCREAMING_SNAKE_CASE__			,							SCREAMING_SNAKE_CASE__			,							SCREAMING_SNAKE_CASE__				:				List[Any]			=		parser.parse_args_into_dataclasses()
     configure_logger(lowercase__     ,  lowercase__ )

     # Downloading and loading a dataset from the hub.
     SCREAMING_SNAKE_CASE__				:				Optional[int]			=		load_dataset(data_args.dataset_name     ,  data_args.dataset_config_name     ,  cache_dir=model_args.cache_dir )

     if "validation" not in datasets.keys():
          # make sure only "validation" and "train" keys remain"
          SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		DatasetDict()
          SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		load_dataset(
              data_args.dataset_name     ,  data_args.dataset_config_name     ,  split=f'''{data_args.train_split_name}[:{data_args.validation_split_percentage}%]'''     ,  cache_dir=model_args.cache_dir     ,  )
          SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		load_dataset(
              data_args.dataset_name     ,  data_args.dataset_config_name     ,  split=f'''{data_args.train_split_name}[{data_args.validation_split_percentage}%:]'''     ,  cache_dir=model_args.cache_dir     ,  )
     else:
          # make sure only "validation" and "train" keys remain"
          SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		DatasetDict()
          SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		load_dataset(
              data_args.dataset_name     ,  data_args.dataset_config_name     ,  split='validation'     ,  cache_dir=model_args.cache_dir     ,  )
          SCREAMING_SNAKE_CASE__				:				Any			=		load_dataset(
              data_args.dataset_name     ,  data_args.dataset_config_name     ,  split=f'''{data_args.train_split_name}'''     ,  cache_dir=model_args.cache_dir     ,  )

     # only normalized-inputs-training is supported
     SCREAMING_SNAKE_CASE__				:				Dict			=		WavaVecaFeatureExtractor.from_pretrained(
         model_args.model_name_or_path     ,  cache_dir=model_args.cache_dir     ,  do_normalize=lowercase__ )

     def prepare_dataset(lowercase__  :			List[Any] ):
          # check that all files have the correct sampling rate
          SCREAMING_SNAKE_CASE__			,							SCREAMING_SNAKE_CASE__				:				Dict			=		librosa.load(batch[data_args.speech_file_column]     ,  sr=feature_extractor.sampling_rate )
          return batch

     # load audio files into numpy arrays
     SCREAMING_SNAKE_CASE__				:				List[str]			=		datasets.map(
         lowercase__     ,  num_proc=data_args.preprocessing_num_workers     ,  remove_columns=datasets['train'].column_names )

     # filter audio files that are too long
     SCREAMING_SNAKE_CASE__				:				Dict			=		vectorized_datasets.filter(
         lambda lowercase__ : len(data['speech'] ) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate ) )

     def normalize(lowercase__  :			List[Any] ):
          return feature_extractor(batch['speech']     ,  sampling_rate=feature_extractor.sampling_rate )

     # normalize and transform to `BatchFeatures`
     SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		vectorized_datasets.map(
         lowercase__     ,  batched=lowercase__     ,  num_proc=data_args.preprocessing_num_workers     ,  load_from_cache_file=not data_args.overwrite_cache     ,  remove_columns=vectorized_datasets['train'].column_names     ,  )

     # pretraining is only supported for "newer" stable layer norm architecture
     # apply_spec_augment has to be True, mask_feature_prob has to be 0.0
     SCREAMING_SNAKE_CASE__				:				Dict			=		WavaVecaConfig.from_pretrained(
         model_args.model_name_or_path     ,  cache_dir=model_args.cache_dir     ,  gradient_checkpointing=training_args.gradient_checkpointing     ,  )

     if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
          raise ValueError(
              'PreTraining is only supported for ``config.do_stable_layer_norm=True`` and'
              ' ``config.feat_extract_norm=\'layer\'' )

     SCREAMING_SNAKE_CASE__				:				int			=		WavaVecaForPreTraining(lowercase__ )

     SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		DataCollatorForWavaVecaPretraining(model=lowercase__     ,  feature_extractor=lowercase__ )

     SCREAMING_SNAKE_CASE__				:				Dict			=		WavaVecaPreTrainer(
         model=lowercase__     ,  data_collator=lowercase__     ,  args=lowercase__     ,  train_dataset=vectorized_datasets['train']     ,  eval_dataset=vectorized_datasets['validation']     ,  tokenizer=lowercase__     ,  max_gumbel_temp=model_args.max_gumbel_temperature     ,  min_gumbel_temp=model_args.min_gumbel_temperature     ,  gumbel_temp_decay=model_args.gumbel_temperature_decay     ,  )
     trainer.train()


if __name__ == "__main__":
      main()




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()
SCREAMING_SNAKE_CASE__		:					int						=					logging.get_logger(__name__)







def     _a				(      lowercase__  :			Union[str, Any]     ,  lowercase__  :			Union[str, Any]=False     ,  lowercase__  :			str=False     ,  lowercase__  :			Dict=False ):







     '''simple docstring'''


     SCREAMING_SNAKE_CASE__				:				Optional[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     _a				(      lowercase__  :			List[str]     ,  lowercase__  :			Dict ):







     '''simple docstring'''


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







def     _a				(      lowercase__  :			Optional[int] ):







     '''simple docstring'''


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







def     _a				(      lowercase__  :			int     ,  lowercase__  :			int     ,  lowercase__  :			Dict ):







     '''simple docstring'''


     SCREAMING_SNAKE_CASE__				:				Tuple			=		dct.pop(lowercase__ )
     SCREAMING_SNAKE_CASE__				:				Any			=		val







@torch.no_grad()
def     _a				(      lowercase__  :			Dict     ,  lowercase__  :			Optional[Any] ):







     '''simple docstring'''


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

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

     # Define processor
     SCREAMING_SNAKE_CASE__				:				str			=		ViltImageProcessor(size=3_84 )
     SCREAMING_SNAKE_CASE__				:				List[Any]			=		BertTokenizer.from_pretrained('bert-base-uncased' )
     SCREAMING_SNAKE_CASE__				:				List[Any]			=		ViltProcessor(lowercase__     ,  lowercase__ )

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

     # Verify outputs
     if mlm_model:
          SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		torch.Size([1, 11, 3_05_22] )
          SCREAMING_SNAKE_CASE__				:				List[str]			=		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"
          SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		outputs.logits[0, 4, :].argmax(-1 ).item()
          assert tokenizer.decode([predicted_id] ) == "cats"
     elif vqa_model:
          SCREAMING_SNAKE_CASE__				:				str			=		torch.Size([1, 31_29] )
          SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		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"
          SCREAMING_SNAKE_CASE__				:				List[Any]			=		outputs.logits.argmax(-1 ).item()
          assert model.config.idalabel[predicted_idx] == "2"
     elif nlvr_model:
          SCREAMING_SNAKE_CASE__				:				Optional[int]			=		torch.Size([1, 2] )
          SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		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__":
      SCREAMING_SNAKE_CASE__		:					str						=					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."
      )

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



