
"""simple docstring"""




import logging
import os
import sys
from dataclasses import dataclass, field
from importlib import import_module
from typing import Dict, List, Optional, Tuple

import numpy as np
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch import nn
from utils_ner import Split, TokenClassificationDataset, TokenClassificationTask

import transformers
from transformers import (
    AutoConfig,
    AutoModelForTokenClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import is_main_process


A__   :					Tuple							=   logging.getLogger(__name__)


@dataclass
class 			__magic_name__       :
       UpperCamelCase_   =		field(
           metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''} )
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''} )
       UpperCamelCase_   =		field(
           default='''NER'''   ,					metadata={'''help''': '''Task type to fine tune in training (e.g. NER, POS, etc)'''} )
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Pretrained tokenizer name or path if not the same as model_name'''} )
       UpperCamelCase_   =		field(default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Set this flag to use fast tokenization.'''} )
       # If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
       # or just modify its tokenizer_config.json.
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}   ,					)


@dataclass
class 			__magic_name__       :
       UpperCamelCase_   =		field(
           metadata={'''help''': '''The input data dir. Should contain the .txt files for a CoNLL-2003-formatted task.'''} )
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Path to a file containing all labels. If not specified, CoNLL-2003 labels are used.'''}   ,					)
       UpperCamelCase_   =		field(
           default=128   ,					metadata={
               '''help''': (
                   '''The maximum total input sequence length after tokenization. Sequences longer '''
                   '''than this will be truncated, sequences shorter will be padded.'''
               )
           }   ,					)
       UpperCamelCase_   =		field(
           default=SCREAMING_SNAKE_CASE__   ,					metadata={'''help''': '''Overwrite the cached training and evaluation sets'''} )




def  _lowerCAmelCase		(		):






  """simple docstring"""



  _lowercase:			int       =     HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)
  if len(sys.argv						) == 2 and sys.argv[1].endswith('''.json'''						):
    # If we pass only one argument to the script and it's the path to a json file,
    # let's parse it to get our arguments.
    _lowercase ,				_lowercase ,				_lowercase:			Tuple       =     parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]						)						)
  else:
    _lowercase ,				_lowercase ,				_lowercase:			Any       =     parser.parse_args_into_dataclasses()

  if (
      os.path.exists(training_args.output_dir						)
      and os.listdir(training_args.output_dir						)
      and training_args.do_train
      and not training_args.overwrite_output_dir
  ):
    raise ValueError(
        f'''Output directory ({training_args.output_dir}) already exists and is not empty. Use'''
        ''' --overwrite_output_dir to overcome.'''						)

  _lowercase:			Union[str, Any]       =     import_module('''tasks'''						)
  try:
    _lowercase:			Union[str, Any]       =     getattr(_UpperCamelCase ,				model_args.task_type						)
    _lowercase:			Tuple       =     token_classification_task_clazz()
  except AttributeError:
    raise ValueError(
        f'''Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. '''
        f'''Available tasks classes are: {TokenClassificationTask.__subclasses__()}'''						)

  # Setup logging
  logging.basicConfig(
      format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''' ,				datefmt='''%m/%d/%Y %H:%M:%S''' ,				level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN ,				)
  logger.warning(
      '''Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s''' ,				training_args.local_rank ,				training_args.device ,				training_args.n_gpu ,				bool(training_args.local_rank != -1						) ,				training_args.fpaa ,				)
  # Set the verbosity to info of the Transformers logger (on main process only):
  if is_main_process(training_args.local_rank						):
    transformers.utils.logging.set_verbosity_info()
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()
  logger.info('''Training/evaluation parameters %s''' ,				_UpperCamelCase						)

  # Set seed
  set_seed(training_args.seed						)

  # Prepare CONLL-2003 task
  _lowercase:			Dict       =     token_classification_task.get_labels(data_args.labels						)
  _lowercase:			Tuple       =     dict(enumerate(_UpperCamelCase						)						)
  _lowercase:			Any       =     len(_UpperCamelCase						)

  # Load pretrained model and tokenizer
  #
  # Distributed training:
  # The .from_pretrained methods guarantee that only one local process can concurrently
  # download model & vocab.

  _lowercase:			Any       =     AutoConfig.from_pretrained(
      model_args.config_name if model_args.config_name else model_args.model_name_or_path ,				num_labels=_UpperCamelCase ,				idalabel=_UpperCamelCase ,				labelaid={label: i for i, label in enumerate(_UpperCamelCase						)} ,				cache_dir=model_args.cache_dir ,				)
  _lowercase:			Union[str, Any]       =     AutoTokenizer.from_pretrained(
      model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path ,				cache_dir=model_args.cache_dir ,				use_fast=model_args.use_fast ,				)
  _lowercase:			Any       =     AutoModelForTokenClassification.from_pretrained(
      model_args.model_name_or_path ,				from_tf=bool('''.ckpt''' in model_args.model_name_or_path						) ,				config=_UpperCamelCase ,				cache_dir=model_args.cache_dir ,				)

  # Get datasets
  _lowercase:			List[str]       =     (
      TokenClassificationDataset(
          token_classification_task=_UpperCamelCase ,				data_dir=data_args.data_dir ,				tokenizer=_UpperCamelCase ,				labels=_UpperCamelCase ,				model_type=config.model_type ,				max_seq_length=data_args.max_seq_length ,				overwrite_cache=data_args.overwrite_cache ,				mode=Split.train ,				)
      if training_args.do_train
      else None
  )
  _lowercase:			Union[str, Any]       =     (
      TokenClassificationDataset(
          token_classification_task=_UpperCamelCase ,				data_dir=data_args.data_dir ,				tokenizer=_UpperCamelCase ,				labels=_UpperCamelCase ,				model_type=config.model_type ,				max_seq_length=data_args.max_seq_length ,				overwrite_cache=data_args.overwrite_cache ,				mode=Split.dev ,				)
      if training_args.do_eval
      else None
  )

  def align_predictions(_UpperCamelCase ,				_UpperCamelCase						) -> Tuple[List[int], List[int]]:
    _lowercase:			Optional[Any]       =     np.argmax(_UpperCamelCase ,				axis=2						)

    _lowercase ,				_lowercase:			int       =     preds.shape

    _lowercase:			str       =     [[] for _ in range(_UpperCamelCase						)]
    _lowercase:			Union[str, Any]       =     [[] for _ in range(_UpperCamelCase						)]

    for i in range(_UpperCamelCase						):
      for j in range(_UpperCamelCase						):
        if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index:
          out_label_list[i].append(label_map[label_ids[i][j]]						)
          preds_list[i].append(label_map[preds[i][j]]						)

    return preds_list, out_label_list

  def compute_metrics(_UpperCamelCase						) -> Dict:
    _lowercase ,				_lowercase:			Dict       =     align_predictions(p.predictions ,				p.label_ids						)
    return {
        "accuracy_score": accuracy_score(_UpperCamelCase ,				_UpperCamelCase						),
        "precision": precision_score(_UpperCamelCase ,				_UpperCamelCase						),
        "recall": recall_score(_UpperCamelCase ,				_UpperCamelCase						),
        "f1": fa_score(_UpperCamelCase ,				_UpperCamelCase						),
    }

  # Data collator
  _lowercase:			Union[str, Any]       =     DataCollatorWithPadding(_UpperCamelCase ,				pad_to_multiple_of=8						) if training_args.fpaa else None

  # Initialize our Trainer
  _lowercase:			Any       =     Trainer(
      model=_UpperCamelCase ,				args=_UpperCamelCase ,				train_dataset=_UpperCamelCase ,				eval_dataset=_UpperCamelCase ,				compute_metrics=_UpperCamelCase ,				data_collator=_UpperCamelCase ,				)

  # Training
  if training_args.do_train:
    trainer.train(
        model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path						) else None						)
    trainer.save_model()
    # For convenience, we also re-save the tokenizer to the same directory,
    # so that you can share your model easily on huggingface.co/models =)
    if trainer.is_world_process_zero():
      tokenizer.save_pretrained(training_args.output_dir						)

    # Evaluation
  _lowercase:			Optional[int]       =     {}
  if training_args.do_eval:
    logger.info('''*** Evaluate ***'''						)

    _lowercase:			Dict       =     trainer.evaluate()

    _lowercase:			Any       =     os.path.join(training_args.output_dir ,				'''eval_results.txt'''						)
    if trainer.is_world_process_zero():
      with open(_UpperCamelCase ,				'''w'''						) as writer:
        logger.info('''***** Eval results *****'''						)
        for key, value in result.items():
          logger.info('''  %s = %s''' ,				_UpperCamelCase ,				_UpperCamelCase						)
          writer.write('''%s = %s\n''' % (key, value)						)

      results.update(_UpperCamelCase						)

    # Predict
  if training_args.do_predict:
    _lowercase:			Any       =     TokenClassificationDataset(
        token_classification_task=_UpperCamelCase ,				data_dir=data_args.data_dir ,				tokenizer=_UpperCamelCase ,				labels=_UpperCamelCase ,				model_type=config.model_type ,				max_seq_length=data_args.max_seq_length ,				overwrite_cache=data_args.overwrite_cache ,				mode=Split.test ,				)

    _lowercase ,				_lowercase ,				_lowercase:			Tuple       =     trainer.predict(_UpperCamelCase						)
    _lowercase ,				_lowercase:			List[str]       =     align_predictions(_UpperCamelCase ,				_UpperCamelCase						)

    _lowercase:			Optional[int]       =     os.path.join(training_args.output_dir ,				'''test_results.txt'''						)
    if trainer.is_world_process_zero():
      with open(_UpperCamelCase ,				'''w'''						) as writer:
        for key, value in metrics.items():
          logger.info('''  %s = %s''' ,				_UpperCamelCase ,				_UpperCamelCase						)
          writer.write('''%s = %s\n''' % (key, value)						)

        # Save predictions
    _lowercase:			List[str]       =     os.path.join(training_args.output_dir ,				'''test_predictions.txt'''						)
    if trainer.is_world_process_zero():
      with open(_UpperCamelCase ,				'''w'''						) as writer:
        with open(os.path.join(data_args.data_dir ,				'''test.txt'''						) ,				'''r'''						) as f:
          token_classification_task.write_predictions_to_file(_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase						)

  return results




def  _lowerCAmelCase		(		_UpperCamelCase						):






  """simple docstring"""



  main()


if __name__ == "__main__":
    main()


"""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()
A__   :					str							=   logging.get_logger(__name__)




def  _lowerCAmelCase		(		_UpperCamelCase ,				_UpperCamelCase=False ,				_UpperCamelCase=False ,				_UpperCamelCase=False						):






  """simple docstring"""



  _lowercase:			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  _lowerCAmelCase		(		_UpperCamelCase ,				_UpperCamelCase						):






  """simple docstring"""



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




def  _lowerCAmelCase		(		_UpperCamelCase						):






  """simple docstring"""



  _lowercase:			Optional[Any]       =     ['''head.weight''', '''head.bias''']
  for k in ignore_keys:
    state_dict.pop(_UpperCamelCase ,				_UpperCamelCase						)




def  _lowerCAmelCase		(		_UpperCamelCase ,				_UpperCamelCase ,				_UpperCamelCase						):






  """simple docstring"""



  _lowercase:			Tuple       =     dct.pop(_UpperCamelCase						)
  _lowercase:			Optional[Any]       =     val




@torch.no_grad()
def  _lowerCAmelCase		(		_UpperCamelCase ,				_UpperCamelCase						):






  """simple docstring"""



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

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

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

  # Define processor
  _lowercase:			List[str]       =     ViltImageProcessor(size=384						)
  _lowercase:			Tuple       =     BertTokenizer.from_pretrained('''bert-base-uncased'''						)
  _lowercase:			Tuple       =     ViltProcessor(_UpperCamelCase ,				_UpperCamelCase						)

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

  # Verify outputs
  if mlm_model:
    _lowercase:			List[str]       =     torch.Size([1, 11, 30_522]						)
    _lowercase:			str       =     torch.tensor([-12.5_061, -12.5_123, -12.5_174]						)
    assert outputs.logits.shape == expected_shape
    assert torch.allclose(outputs.logits[0, 0, :3] ,				_UpperCamelCase ,				atol=1e-4						)

    # verify masked token prediction equals "cats"
    _lowercase:			Union[str, Any]       =     outputs.logits[0, 4, :].argmax(-1						).item()
    assert tokenizer.decode([predicted_id]						) == "cats"
  elif vqa_model:
    _lowercase:			int       =     torch.Size([1, 3_129]						)
    _lowercase:			Optional[Any]       =     torch.tensor([-15.9_495, -18.1_472, -10.3_041]						)
    assert torch.allclose(outputs.logits[0, :3] ,				_UpperCamelCase ,				atol=1e-4						)
    assert outputs.logits.shape == expected_shape
    assert torch.allclose(outputs.logits[0, 0, :3] ,				_UpperCamelCase ,				atol=1e-4						)

    # verify vqa prediction equals "2"
    _lowercase:			Optional[Any]       =     outputs.logits.argmax(-1						).item()
    assert model.config.idalabel[predicted_idx] == "2"
  elif nlvr_model:
    _lowercase:			List[Any]       =     torch.Size([1, 2]						)
    _lowercase:			int       =     torch.tensor([-2.8_721, 2.1_291]						)
    assert torch.allclose(outputs.logits[0, :3] ,				_UpperCamelCase ,				atol=1e-4						)
    assert outputs.logits.shape == expected_shape

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


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

  A__   :					Union[str, Any]							=   parser.parse_args()
  convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
