'''simple docstring'''


import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
					import os

					import jax.numpy as jnp
					from jax import jit

					from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
					from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

					__lowerCAmelCase						:							Optional[int]    							=						"0.12"  # assumed parallelism: 8


if is_torch_available():
					import torch





def    lowerCAmelCase       (		UpperCamelCase__					: Union[str, Any]     ,       UpperCamelCase__					: Any     ,       UpperCamelCase__					: Optional[int]=None							):


	"""simple docstring"""
	if rng is None:
		__UpperCAmelCase      					=      random.Random()

	__UpperCAmelCase      					=      1
	for dim in shape:
		total_dims *= dim

	__UpperCAmelCase      					=      []
	for _ in range(UpperCamelCase__							):
		values.append(rng.randint(0     ,       vocab_size - 1							)							)

	__UpperCAmelCase      					=      np.array(UpperCamelCase__     ,       dtype=jnp.intaa							).reshape(UpperCamelCase__							)

	return output





def    lowerCAmelCase       (		UpperCamelCase__					: Union[str, Any]     ,       UpperCamelCase__					: Optional[Any]=None							):


	"""simple docstring"""
	__UpperCAmelCase      					=      ids_tensor(UpperCamelCase__     ,       vocab_size=2     ,       rng=UpperCamelCase__							)
	# make sure that at least one token is attended to for each batch
	__UpperCAmelCase      					=      1
	return attn_mask






@require_flax
class     A :
							a_										=   None
							a_										=   ()





							def    snake_case__    (   self					:				List[Any]						)				->						Any:
								__UpperCAmelCase					,      __UpperCAmelCase      					=      self.model_tester.prepare_config_and_inputs_for_common()

								# cut to half length & take max batch_size 3
								__UpperCAmelCase      					=      2
								__UpperCAmelCase      					=      inputs['''input_ids'''].shape[-1] // 2
								__UpperCAmelCase      					=      inputs['''input_ids'''][:max_batch_size, :sequence_length]

								__UpperCAmelCase      					=      jnp.ones_like(lowerCamelCase_						)
								__UpperCAmelCase      					=      attention_mask[:max_batch_size, :sequence_length]

								# generate max 5 tokens
								__UpperCAmelCase      					=      input_ids.shape[-1] + 5
								if config.eos_token_id is not None and config.pad_token_id is None:
									# hack to allow generate for models such as GPT2 as is done in `generate()`
									__UpperCAmelCase      					=      config.eos_token_id
								return config, input_ids, attention_mask, max_length





							@is_pt_flax_cross_test
							def    snake_case__    (   self					:				List[str]						)				->						List[str]:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      False
								__UpperCAmelCase      					=      max_length
								__UpperCAmelCase      					=      0

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model_class.__name__[4:]  # Skip the "Flax" at the beginning
									__UpperCAmelCase      					=      getattr(lowerCamelCase_  ,		lowerCamelCase_						)
									__UpperCAmelCase      					=      pt_model_class(lowerCamelCase_						).eval()
									__UpperCAmelCase      					=      load_flax_weights_in_pytorch_model(lowerCamelCase_  ,		flax_model.params						)

									__UpperCAmelCase      					=      flax_model.generate(lowerCamelCase_						).sequences
									__UpperCAmelCase      					=      pt_model.generate(torch.tensor(lowerCamelCase_  ,		dtype=torch.long						)						)

									if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
										__UpperCAmelCase      					=      flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

									self.assertListEqual(pt_generation_outputs.numpy().tolist()  ,		flax_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				Any						)				->						str:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      False
								__UpperCAmelCase      					=      max_length

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				Optional[Any]						)				->						Any:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      True
								__UpperCAmelCase      					=      max_length

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				str						)				->						int:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      False
								__UpperCAmelCase      					=      max_length
								__UpperCAmelCase      					=      2

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				Dict						)				->						List[str]:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      False
								__UpperCAmelCase      					=      max_length
								__UpperCAmelCase      					=      2
								__UpperCAmelCase      					=      2

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[0]  ,		input_ids.shape[0] * config.num_return_sequences						)





							def    snake_case__    (   self					:				Union[str, Any]						)				->						int:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      True
								__UpperCAmelCase      					=      max_length
								__UpperCAmelCase      					=      0.8
								__UpperCAmelCase      					=      1_0
								__UpperCAmelCase      					=      0.3
								__UpperCAmelCase      					=      1
								__UpperCAmelCase      					=      8
								__UpperCAmelCase      					=      9

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				List[str]						)				->						str:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      max_length
								__UpperCAmelCase      					=      1
								__UpperCAmelCase      					=      8
								__UpperCAmelCase      					=      9

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				Optional[int]						)				->						Optional[Any]:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()
								__UpperCAmelCase      					=      max_length
								__UpperCAmelCase      					=      2
								__UpperCAmelCase      					=      1
								__UpperCAmelCase      					=      8
								__UpperCAmelCase      					=      9

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				str						)				->						Dict:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()

								# pad attention mask on the left
								__UpperCAmelCase      					=      attention_mask.at[(0, 0)].set(0						)

								__UpperCAmelCase      					=      False
								__UpperCAmelCase      					=      max_length

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_  ,		attention_mask=lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_  ,		attention_mask=lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)





							def    snake_case__    (   self					:				Union[str, Any]						)				->						Dict:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()

								# pad attention mask on the left
								__UpperCAmelCase      					=      attention_mask.at[(0, 0)].set(0						)

								__UpperCAmelCase      					=      True
								__UpperCAmelCase      					=      max_length

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_  ,		attention_mask=lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_  ,		attention_mask=lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)






							def    snake_case__    (   self					:				Optional[Any]						)				->						Tuple:
								__UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase					,      __UpperCAmelCase      					=      self._get_input_ids_and_config()

								# pad attention mask on the left
								__UpperCAmelCase      					=      attention_mask.at[(0, 0)].set(0						)

								__UpperCAmelCase      					=      2
								__UpperCAmelCase      					=      max_length

								for model_class in self.all_generative_model_classes:
									__UpperCAmelCase      					=      model_class(lowerCamelCase_						)

									__UpperCAmelCase      					=      model.generate(lowerCamelCase_  ,		attention_mask=lowerCamelCase_						).sequences
									self.assertEqual(generation_outputs.shape[-1]  ,		lowerCamelCase_						)

									__UpperCAmelCase      					=      jit(model.generate						)
									__UpperCAmelCase      					=      jit_generate(lowerCamelCase_  ,		attention_mask=lowerCamelCase_						).sequences

									self.assertListEqual(generation_outputs.tolist()  ,		jit_generation_outputs.tolist()						)






@require_flax
class     A ( unittest.TestCase ):





							def    snake_case__    (   self					:				str						)				->						Optional[int]:
								__UpperCAmelCase      					=      AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''						)
								__UpperCAmelCase      					=      FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''						)

								__UpperCAmelCase      					=      '''Hello world'''
								__UpperCAmelCase      					=      tokenizer(lowerCamelCase_  ,		return_tensors='''np'''						).input_ids

								# typos are quickly detected (the correct argument is `do_sample`)
								with self.assertRaisesRegex(lowerCamelCase_  ,		'''do_samples'''						):
									model.generate(lowerCamelCase_  ,		do_samples=lowerCamelCase_						)

								# arbitrary arguments that will not be used anywhere are also not accepted
								with self.assertRaisesRegex(lowerCamelCase_  ,		'''foo'''						):
									__UpperCAmelCase      					=      {'''foo''': '''bar'''}
									model.generate(lowerCamelCase_  ,		**lowerCamelCase_						)




'''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()
__UpperCAmelCase       							=				logging.get_logger(__name__)


def    _snake_case  (	A					,       A=False					,       A=False					,       A=False			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      []
							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    _snake_case  (	A					,       A			)		->       List[str]:
							for i in range(config.num_hidden_layers			):
														lowerCAmelCase__			  =      '''vilt.'''
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														lowerCAmelCase__			  =      state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.weight"""			)
														lowerCAmelCase__			  =      state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.bias"""			)
														# next, add query, keys and values (in that order) to the state dict
														lowerCAmelCase__			  =      in_proj_weight[
														    : config.hidden_size, :
														]
														lowerCAmelCase__			  =      in_proj_bias[: config.hidden_size]
														lowerCAmelCase__			  =      in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														lowerCAmelCase__			  =      in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														lowerCAmelCase__			  =      in_proj_weight[
														    -config.hidden_size :, :
														]
														lowerCAmelCase__			  =      in_proj_bias[-config.hidden_size :]


def    _snake_case  (	A			)		->       List[str]:
							lowerCAmelCase__			  =      ['''head.weight''', '''head.bias''']
							for k in ignore_keys:
														state_dict.pop(A					,       A			)


def    _snake_case  (	A					,       A					,       A			)		->       str:
							lowerCAmelCase__			  =      dct.pop(A			)
							lowerCAmelCase__			  =      val


@torch.no_grad()
def    _snake_case  (	A					,       A			)		->       Any:
							lowerCAmelCase__			  =      ViltConfig(image_size=384					,       patch_size=32					,       tie_word_embeddings=A			)
							lowerCAmelCase__			  =      False
							lowerCAmelCase__			  =      False
							lowerCAmelCase__			  =      False
							lowerCAmelCase__			  =      False
							if "vqa" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      3129
														lowerCAmelCase__			  =      '''huggingface/label-files'''
														lowerCAmelCase__			  =      '''vqa2-id2label.json'''
														lowerCAmelCase__			  =      json.load(open(hf_hub_download(A					,       A					,       repo_type='''dataset'''			)					,       '''r'''			)			)
														lowerCAmelCase__			  =      {int(A			): v for k, v in idalabel.items()}
														lowerCAmelCase__			  =      idalabel
														lowerCAmelCase__			  =      {v: k for k, v in idalabel.items()}
														lowerCAmelCase__			  =      ViltForQuestionAnswering(A			)
							elif "nlvr" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      2
														lowerCAmelCase__			  =      {0: '''False''', 1: '''True'''}
														lowerCAmelCase__			  =      {v: k for k, v in config.idalabel.items()}
														lowerCAmelCase__			  =      3
														lowerCAmelCase__			  =      ViltForImagesAndTextClassification(A			)
							elif "irtr" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      ViltForImageAndTextRetrieval(A			)
							elif "mlm_itm" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      ViltForMaskedLM(A			)
							else:
														raise ValueError('''Unknown model type'''			)

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

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

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

							# Verify outputs
							if mlm_model:
														lowerCAmelCase__			  =      torch.Size([1, 11, 30522]			)
														lowerCAmelCase__			  =      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]					,       A					,       atol=1E-4			)

														# verify masked token prediction equals "cats"
														lowerCAmelCase__			  =      outputs.logits[0, 4, :].argmax(-1			).item()
														assert tokenizer.decode([predicted_id]			) == "cats"
							elif vqa_model:
														lowerCAmelCase__			  =      torch.Size([1, 3129]			)
														lowerCAmelCase__			  =      torch.tensor([-15.9_495, -18.1_472, -10.3_041]			)
														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"
														lowerCAmelCase__			  =      outputs.logits.argmax(-1			).item()
														assert model.config.idalabel[predicted_idx] == "2"
							elif nlvr_model:
														lowerCAmelCase__			  =      torch.Size([1, 2]			)
														lowerCAmelCase__			  =      torch.tensor([-2.8_721, 2.1_291]			)
														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__":
		__UpperCAmelCase       							=				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.'''
		)

		__UpperCAmelCase       							=				parser.parse_args()
		convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)