
"""simple docstring"""




import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import UperNetForSemanticSegmentation
  from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor


class       __magic_name__							:
       def __init__(  self			,					A_			,					A_=13			,					A_=32			,					A_=3			,					A_=4			,					A_=[10, 20, 30, 40]			,					A_=[2, 2, 3, 2]			,					A_=True			,					A_=True			,					A_=37			,					A_="gelu"			,					A_=10			,					A_=0.02			,					A_=["stage2", "stage3", "stage4"]			,					A_=3			,					A_=None			,					)  ->						List[str]:





         """simple docstring"""




         _lowercase:      str   			=  parent
         _lowercase:      Union[str, Any]   			=  batch_size
         _lowercase:      Tuple   			=  image_size
         _lowercase:      Tuple   			=  num_channels
         _lowercase:      Optional[int]   			=  num_stages
         _lowercase:      List[Any]   			=  hidden_sizes
         _lowercase:      Dict   			=  depths
         _lowercase:      int   			=  is_training
         _lowercase:      Tuple   			=  use_labels
         _lowercase:      Any   			=  intermediate_size
         _lowercase:      Union[str, Any]   			=  hidden_act
         _lowercase:      str   			=  type_sequence_label_size
         _lowercase:      int   			=  initializer_range
         _lowercase:      List[str]   			=  out_features
         _lowercase:      Any   			=  num_labels
         _lowercase:      Any   			=  scope
         _lowercase:      str   			=  num_stages
       def 				lowercase_				(  self							)  ->						str:





         """simple docstring"""




         _lowercase:      Union[str, Any]   			=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

         _lowercase:      Any   			=  None
         if self.use_labels:
           _lowercase:      Dict   			=  ids_tensor([self.batch_size]			,					self.type_sequence_label_size							)

         _lowercase:      Tuple   			=  self.get_config()

         return config, pixel_values, labels
       def 				lowercase_				(  self							)  ->						List[Any]:





         """simple docstring"""




         return ConvNextConfig(
             num_channels=self.num_channels			,					num_stages=self.num_stages			,					hidden_sizes=self.hidden_sizes			,					depths=self.depths			,					is_training=self.is_training			,					intermediate_size=self.intermediate_size			,					hidden_act=self.hidden_act			,					out_features=self.out_features			,					)
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         return UperNetConfig(
             backbone_config=self.get_backbone_config()			,					hidden_size=512			,					pool_scales=[1, 2, 3, 6]			,					use_auxiliary_head=A_			,					auxiliary_loss_weight=0.4			,					auxiliary_in_channels=40			,					auxiliary_channels=256			,					auxiliary_num_convs=1			,					auxiliary_concat_input=A_			,					loss_ignore_index=255			,					num_labels=self.num_labels			,					)
       def 				lowercase_				(  self			,					A_			,					A_			,					A_							)  ->						int:





         """simple docstring"""




         _lowercase:      List[str]   			=  UperNetForSemanticSegmentation(config=A_							)
         model.to(A_							)
         model.eval()
         _lowercase:      int   			=  model(A_							)
         self.parent.assertEqual(
             result.logits.shape			,					(self.batch_size, self.num_labels, self.image_size, self.image_size)							)



       def 				lowercase_				(  self							)  ->						List[str]:





         """simple docstring"""




         _lowercase:      Optional[int]   			=  self.prepare_config_and_inputs()
         (
             (
             _lowercase
         )			,					(
             _lowercase
         )			,					(
             _lowercase
         )			,					
         ):      Tuple   			=  config_and_inputs
         _lowercase:      str   			=  {'''pixel_values''': pixel_values}
         return config, inputs_dict


@require_torch
class       __magic_name__							(							SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__ ,       unittest.TestCase			):
       UpperCamelCase_ =		(UperNetForSemanticSegmentation,) if is_torch_available() else ()
       UpperCamelCase_ =		{'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
       UpperCamelCase_ =		False
       UpperCamelCase_ =		False
       UpperCamelCase_ =		False
       UpperCamelCase_ =		False
       UpperCamelCase_ =		False
       UpperCamelCase_ =		False
       def 				lowercase_				(  self							)  ->						Optional[int]:





         """simple docstring"""




         _lowercase:      List[Any]   			=  UperNetModelTester(self							)
         _lowercase:      Optional[Any]   			=  ConfigTester(self			,					config_class=A_			,					has_text_modality=A_			,					hidden_size=37							)
       def 				lowercase_				(  self							)  ->						List[str]:





         """simple docstring"""




         self.create_and_test_config_common_properties()
         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         return
       def 				lowercase_				(  self							)  ->						Dict:





         """simple docstring"""




         _lowercase			,					_lowercase:      Union[str, Any]   			=  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           _lowercase:      Optional[Any]   			=  model_class(A_							)
           _lowercase:      int   			=  inspect.signature(model.forward							)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           _lowercase:      Dict   			=  [*signature.parameters.keys()]

           _lowercase:      int   			=  ['''pixel_values''']
           self.assertListEqual(arg_names[:1]			,					A_							)
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         _lowercase:      List[str]   			=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_semantic_segmentation(*A_							)
       @unittest.skip(reason='''UperNet does not use inputs_embeds'''							)
       def 				lowercase_				(  self							)  ->						str:





         """simple docstring"""




         pass
       @unittest.skip(reason='''UperNet does not support input and output embeddings'''							)
       def 				lowercase_				(  self							)  ->						List[Any]:





         """simple docstring"""




         pass
       @unittest.skip(reason='''UperNet does not have a base model'''							)
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         pass
       @unittest.skip(reason='''UperNet does not have a base model'''							)
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         pass
       @require_torch_multi_gpu
       @unittest.skip(reason='''UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''							)
       def 				lowercase_				(  self							)  ->						Optional[Any]:





         """simple docstring"""




         pass
       @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''							)
       def 				lowercase_				(  self							)  ->						List[Any]:





         """simple docstring"""




         pass
       def 				lowercase_				(  self							)  ->						str:





         """simple docstring"""




         def check_hidden_states_output(A_			,					A_			,					A_							):
           _lowercase:      Dict   			=  model_class(A_							)
           model.to(A_							)
           model.eval()

           with torch.no_grad():
             _lowercase:      List[str]   			=  model(**self._prepare_for_class(A_			,					A_							)							)

           _lowercase:      Tuple   			=  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

           _lowercase:      Union[str, Any]   			=  self.model_tester.num_stages
           self.assertEqual(len(A_							)			,					expected_num_stages + 1							)

           # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
           self.assertListEqual(
               list(hidden_states[0].shape[-2:]							)			,					[self.model_tester.image_size // 4, self.model_tester.image_size // 4]			,					)

         _lowercase			,					_lowercase:      Dict   			=  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           _lowercase:      Optional[int]   			=  True
           check_hidden_states_output(A_			,					A_			,					A_							)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           _lowercase:      Union[str, Any]   			=  True

           check_hidden_states_output(A_			,					A_			,					A_							)
       def 				lowercase_				(  self							)  ->						Union[str, Any]:





         """simple docstring"""




         _lowercase			,					_lowercase:      Any   			=  self.model_tester.prepare_config_and_inputs_for_common()

         _lowercase:      int   			=  _config_zero_init(A_							)
         _lowercase:      Optional[Any]   			=  _config_zero_init(configs_no_init.backbone_config							)
         for model_class in self.all_model_classes:
           _lowercase:      Any   			=  model_class(config=A_							)
           for name, param in model.named_parameters():
             if param.requires_grad:
               self.assertIn(
                   ((param.data.mean() * 1E9).round() / 1E9).item()			,					[0.0, 1.0]			,					msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''			,					)
       @unittest.skip(reason='''UperNet does not have tied weights'''							)
       def 				lowercase_				(  self							)  ->						List[str]:





         """simple docstring"""




         pass



       @slow
       def 				lowercase_				(  self							)  ->						Tuple:





         """simple docstring"""




         for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _lowercase:      Union[str, Any]   			=  UperNetForSemanticSegmentation.from_pretrained(A_							)
           self.assertIsNotNone(A_							)




def 			_lowerCAmelCase		(		):






  """simple docstring"""



  _lowercase:      List[str]   			=  hf_hub_download(
      repo_id='''hf-internal-testing/fixtures_ade20k'''			,					repo_type='''dataset'''			,					filename='''ADE_val_00000001.jpg'''    )
  _lowercase:      Optional[int]   			=  Image.open(_UpperCamelCase    ).convert('''RGB'''    )
  return image







@require_torch
@require_vision
@slow
class       __magic_name__							(							unittest.TestCase			):
       def 				lowercase_				(  self							)  ->						int:





         """simple docstring"""




         _lowercase:      Optional[Any]   			=  AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''							)
         _lowercase:      Optional[Any]   			=  UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''							).to(A_							)

         _lowercase:      List[Any]   			=  prepare_img()
         _lowercase:      Any   			=  processor(images=A_			,					return_tensors='''pt'''							).to(A_							)

         with torch.no_grad():
           _lowercase:      str   			=  model(**A_							)

         _lowercase:      Union[str, Any]   			=  torch.Size((1, model.config.num_labels, 512, 512)							)
         self.assertEqual(outputs.logits.shape			,					A_							)

         _lowercase:      Optional[Any]   			=  torch.tensor(
             [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]							).to(A_							)
         self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]			,					A_			,					atol=1E-4							)							)



       def 				lowercase_				(  self							)  ->						int:





         """simple docstring"""




         _lowercase:      Tuple   			=  AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''							)
         _lowercase:      str   			=  UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''							).to(A_							)

         _lowercase:      Union[str, Any]   			=  prepare_img()
         _lowercase:      Optional[Any]   			=  processor(images=A_			,					return_tensors='''pt'''							).to(A_							)

         with torch.no_grad():
           _lowercase:      str   			=  model(**A_							)

         _lowercase:      Optional[Any]   			=  torch.Size((1, model.config.num_labels, 512, 512)							)
         self.assertEqual(outputs.logits.shape			,					A_							)

         _lowercase:      Tuple   			=  torch.tensor(
             [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]							).to(A_							)
         self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]			,					A_			,					atol=1E-4							)							)


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




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






  """simple docstring"""



  _lowercase:      Tuple   			=  []
  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:      List[str]   			=  '''vilt.'''
    # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
    _lowercase:      Optional[int]   			=  state_dict.pop(f'''transformer.blocks.{i}.attn.qkv.weight'''    )
    _lowercase:      int   			=  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:      List[Any]   			=  in_proj_bias[: config.hidden_size]
    _lowercase:      Dict   			=  in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    _lowercase:      List[Any]   			=  in_proj_bias[
        config.hidden_size : config.hidden_size * 2
    ]
    _lowercase:      Optional[Any]   			=  in_proj_weight[
        -config.hidden_size :, :
    ]
    _lowercase:      Optional[Any]   			=  in_proj_bias[-config.hidden_size :]




def 			_lowerCAmelCase		(		_UpperCamelCase    ):






  """simple docstring"""



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




def 			_lowerCAmelCase		(		_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase    ):






  """simple docstring"""



  _lowercase:      List[str]   			=  dct.pop(_UpperCamelCase    )
  _lowercase:      Dict   			=  val




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






  """simple docstring"""



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

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

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

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

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

  A__ :       str				= parser.parse_args()
  convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
