


'''simple docstring'''






import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.generation import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)




@require_flax
class 						a__							(				unittest.TestCase    ):





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Dict:
														lowerCAmelCase__			  =      jnp.ones((batch_size, length)  ) / length
														return scores







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      None
														lowerCAmelCase__			  =      20

														lowerCAmelCase__			  =      self._get_uniform_logits(batch_size=2   ,       length=lowerCamelCase_  )

														# tweak scores to not be uniform anymore
														lowerCAmelCase__			  =      scores.at[1, 5].set((1 / length) + 0.1  )  # peak, 1st batch
														lowerCAmelCase__			  =      scores.at[1, 10].set((1 / length) - 0.4  )  # valley, 1st batch

														# compute softmax
														lowerCAmelCase__			  =      jax.nn.softmax(lowerCamelCase_   ,       axis=-1  )

														lowerCAmelCase__			  =      FlaxTemperatureLogitsWarper(temperature=0.5  )
														lowerCAmelCase__			  =      FlaxTemperatureLogitsWarper(temperature=1.3  )

														lowerCAmelCase__			  =      jax.nn.softmax(temp_dist_warper_sharper(lowerCamelCase_   ,       scores.copy()   ,       cur_len=lowerCamelCase_  )   ,       axis=-1  )
														lowerCAmelCase__			  =      jax.nn.softmax(temp_dist_warper_smoother(lowerCamelCase_   ,       scores.copy()   ,       cur_len=lowerCamelCase_  )   ,       axis=-1  )

														# uniform distribution stays uniform
														self.assertTrue(jnp.allclose(probs[0, :]   ,       warped_prob_sharp[0, :]   ,       atol=1e-3  )  )
														self.assertTrue(jnp.allclose(probs[0, :]   ,       warped_prob_smooth[0, :]   ,       atol=1e-3  )  )

														# sharp peaks get higher, valleys get lower
														self.assertLess(probs[1, :].max()   ,       warped_prob_sharp[1, :].max()  )
														self.assertGreater(probs[1, :].min()   ,       warped_prob_sharp[1, :].min()  )

														# smooth peaks get lower, valleys get higher
														self.assertGreater(probs[1, :].max()   ,       warped_prob_smooth[1, :].max()  )
														self.assertLess(probs[1, :].min()   ,       warped_prob_smooth[1, :].min()  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      None
														lowerCAmelCase__			  =      10
														lowerCAmelCase__			  =      2

														# create ramp distribution
														lowerCAmelCase__			  =      np.broadcast_to(np.arange(lowerCamelCase_  )[None, :]   ,       (batch_size, vocab_size)  ).copy()
														lowerCAmelCase__			  =      ramp_logits[1:, : vocab_size // 2] + vocab_size

														lowerCAmelCase__			  =      FlaxTopKLogitsWarper(3  )

														lowerCAmelCase__			  =      top_k_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )

														# check that correct tokens are filtered
														self.assertListEqual(jnp.isinf(scores[0]  ).tolist()   ,       7 * [True] + 3 * [False]  )
														self.assertListEqual(jnp.isinf(scores[1]  ).tolist()   ,       2 * [True] + 3 * [False] + 5 * [True]  )

														# check special case
														lowerCAmelCase__			  =      5
														lowerCAmelCase__			  =      FlaxTopKLogitsWarper(top_k=1   ,       filter_value=0.0   ,       min_tokens_to_keep=3  )

														lowerCAmelCase__			  =      np.broadcast_to(np.arange(lowerCamelCase_  )[None, :]   ,       (batch_size, length)  ).copy()
														lowerCAmelCase__			  =      top_k_warp_safety_check(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )

														# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
														self.assertListEqual((scores == 0.0).sum(axis=-1  ).tolist()   ,       [2, 2]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      None
														lowerCAmelCase__			  =      10
														lowerCAmelCase__			  =      2

														# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
														lowerCAmelCase__			  =      np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]  )  )

														lowerCAmelCase__			  =      FlaxTopPLogitsWarper(0.8  )
														lowerCAmelCase__			  =      np.exp(top_p_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )  )

														# dist should be filtered to keep min num values so that sum is >= top_p
														# exp (-inf) => 0
														lowerCAmelCase__			  =      np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]  )
														self.assertTrue(np.allclose(lowerCamelCase_   ,       lowerCamelCase_   ,       atol=1e-3  )  )

														# check edge cases with negative and extreme logits
														lowerCAmelCase__			  =      np.broadcast_to(np.arange(lowerCamelCase_  )[None, :]   ,       (batch_size, vocab_size)  ).copy() - (
														    vocab_size // 2
														)

														# make ramp_logits more extreme
														lowerCAmelCase__			  =      ramp_logits[1] * 100.0

														# make sure at least 2 tokens are kept
														lowerCAmelCase__			  =      FlaxTopPLogitsWarper(0.9   ,       min_tokens_to_keep=2   ,       filter_value=0.0  )
														lowerCAmelCase__			  =      top_p_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )

														# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
														self.assertListEqual((filtered_dist != 0.0).sum(axis=-1  ).tolist()   ,       [3, 2]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      20
														lowerCAmelCase__			  =      4
														lowerCAmelCase__			  =      0

														lowerCAmelCase__			  =      FlaxMinLengthLogitsProcessor(min_length=10   ,       eos_token_id=lowerCamelCase_  )

														# check that min length is applied at length 5
														lowerCAmelCase__			  =      ids_tensor((batch_size, 20)   ,       vocab_size=20  )
														lowerCAmelCase__			  =      5
														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      min_dist_processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()   ,       4 * [-float('''inf'''  )]  )

														# check that min length is not applied anymore at length 15
														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      15
														lowerCAmelCase__			  =      min_dist_processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														self.assertFalse(jnp.isinf(lowerCamelCase_  ).any()  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														lowerCAmelCase__			  =      20
														lowerCAmelCase__			  =      4
														lowerCAmelCase__			  =      0

														lowerCAmelCase__			  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_  )

														# check that all scores are -inf except the bos_token_id score
														lowerCAmelCase__			  =      ids_tensor((batch_size, 1)   ,       vocab_size=20  )
														lowerCAmelCase__			  =      1
														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      logits_processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]  ).all()  )
														self.assertListEqual(scores[:, bos_token_id].tolist()   ,       4 * [0]  )  # score for bos_token_id shold be zero

														# check that bos_token_id is not forced if current length is greater than 1
														lowerCAmelCase__			  =      3
														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      logits_processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														self.assertFalse(jnp.isinf(lowerCamelCase_  ).any()  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      20
														lowerCAmelCase__			  =      4
														lowerCAmelCase__			  =      0
														lowerCAmelCase__			  =      5

														lowerCAmelCase__			  =      FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_   ,       eos_token_id=lowerCamelCase_  )

														# check that all scores are -inf except the eos_token_id when max_length is reached
														lowerCAmelCase__			  =      ids_tensor((batch_size, 4)   ,       vocab_size=20  )
														lowerCAmelCase__			  =      4
														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      logits_processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]  ).all()  )
														self.assertListEqual(scores[:, eos_token_id].tolist()   ,       4 * [0]  )  # score for eos_token_id should be zero

														# check that eos_token_id is not forced if max_length is not reached
														lowerCAmelCase__			  =      3
														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      logits_processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														self.assertFalse(jnp.isinf(lowerCamelCase_  ).any()  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														lowerCAmelCase__			  =      4
														lowerCAmelCase__			  =      10
														lowerCAmelCase__			  =      15
														lowerCAmelCase__			  =      2
														lowerCAmelCase__			  =      1
														lowerCAmelCase__			  =      15

														# dummy input_ids and scores
														lowerCAmelCase__			  =      ids_tensor((batch_size, sequence_length)   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      input_ids.copy()

														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      scores.copy()

														# instantiate all dist processors
														lowerCAmelCase__			  =      FlaxTemperatureLogitsWarper(temperature=0.5  )
														lowerCAmelCase__			  =      FlaxTopKLogitsWarper(3  )
														lowerCAmelCase__			  =      FlaxTopPLogitsWarper(0.8  )

														# instantiate all logits processors
														lowerCAmelCase__			  =      FlaxMinLengthLogitsProcessor(min_length=10   ,       eos_token_id=lowerCamelCase_  )
														lowerCAmelCase__			  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_  )
														lowerCAmelCase__			  =      FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_   ,       eos_token_id=lowerCamelCase_  )

														lowerCAmelCase__			  =      10

														# no processor list
														lowerCAmelCase__			  =      temp_dist_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														lowerCAmelCase__			  =      top_k_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														lowerCAmelCase__			  =      top_p_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														lowerCAmelCase__			  =      min_dist_proc(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														lowerCAmelCase__			  =      bos_dist_proc(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
														lowerCAmelCase__			  =      eos_dist_proc(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )

														# with processor list
														lowerCAmelCase__			  =      FlaxLogitsProcessorList(
														    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
														lowerCAmelCase__			  =      processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )

														# scores should be equal
														self.assertTrue(jnp.allclose(lowerCamelCase_   ,       lowerCamelCase_   ,       atol=1e-3  )  )

														# input_ids should never be changed
														self.assertListEqual(input_ids.tolist()   ,       input_ids_comp.tolist()  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      4
														lowerCAmelCase__			  =      10
														lowerCAmelCase__			  =      15
														lowerCAmelCase__			  =      2
														lowerCAmelCase__			  =      1
														lowerCAmelCase__			  =      15

														# dummy input_ids and scores
														lowerCAmelCase__			  =      ids_tensor((batch_size, sequence_length)   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      input_ids.copy()

														lowerCAmelCase__			  =      self._get_uniform_logits(lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      scores.copy()

														# instantiate all dist processors
														lowerCAmelCase__			  =      FlaxTemperatureLogitsWarper(temperature=0.5  )
														lowerCAmelCase__			  =      FlaxTopKLogitsWarper(3  )
														lowerCAmelCase__			  =      FlaxTopPLogitsWarper(0.8  )

														# instantiate all logits processors
														lowerCAmelCase__			  =      FlaxMinLengthLogitsProcessor(min_length=10   ,       eos_token_id=lowerCamelCase_  )
														lowerCAmelCase__			  =      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase_  )
														lowerCAmelCase__			  =      FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase_   ,       eos_token_id=lowerCamelCase_  )

														lowerCAmelCase__			  =      10

														# no processor list
														def run_no_processor_list(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  ):
																					lowerCAmelCase__			  =      temp_dist_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
																					lowerCAmelCase__			  =      top_k_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
																					lowerCAmelCase__			  =      top_p_warp(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
																					lowerCAmelCase__			  =      min_dist_proc(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
																					lowerCAmelCase__			  =      bos_dist_proc(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
																					lowerCAmelCase__			  =      eos_dist_proc(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
																					return scores

														# with processor list
														def run_processor_list(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  ):
																					lowerCAmelCase__			  =      FlaxLogitsProcessorList(
																					    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
																					lowerCAmelCase__			  =      processor(lowerCamelCase_   ,       lowerCamelCase_   ,       cur_len=lowerCamelCase_  )
																					return scores

														lowerCAmelCase__			  =      jax.jit(lowerCamelCase_  )
														lowerCAmelCase__			  =      jax.jit(lowerCamelCase_  )

														lowerCAmelCase__			  =      jitted_run_no_processor_list(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )
														lowerCAmelCase__			  =      jitted_run_processor_list(lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_  )

														# scores should be equal
														self.assertTrue(jnp.allclose(lowerCamelCase_   ,       lowerCamelCase_   ,       atol=1e-3  )  )

														# input_ids should never be changed
														self.assertListEqual(input_ids.tolist()   ,       input_ids_comp.tolist()  )



'''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)