
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 :			List[str]              =	logging.get_logger(__name__)


def 				A	( lowercase   ,	lowercase=False   ,	lowercase=False   ,	lowercase=False			)       ->  Optional[int]:
  '''simple docstring'''

  UpperCamelCase          =					[]
  for i in range(config.num_hidden_layers			):
    # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
    rename_keys.append((f'''transformer.blocks.{i}.norm1.weight''', f'''vilt.encoder.layer.{i}.layernorm_before.weight''')			)
    rename_keys.append((f'''transformer.blocks.{i}.norm1.bias''', f'''vilt.encoder.layer.{i}.layernorm_before.bias''')			)
    rename_keys.append(
        (f'''transformer.blocks.{i}.attn.proj.weight''', f'''vilt.encoder.layer.{i}.attention.output.dense.weight''')			)
    rename_keys.append(
        (f'''transformer.blocks.{i}.attn.proj.bias''', f'''vilt.encoder.layer.{i}.attention.output.dense.bias''')			)
    rename_keys.append((f'''transformer.blocks.{i}.norm2.weight''', f'''vilt.encoder.layer.{i}.layernorm_after.weight''')			)
    rename_keys.append((f'''transformer.blocks.{i}.norm2.bias''', f'''vilt.encoder.layer.{i}.layernorm_after.bias''')			)
    rename_keys.append(
        (f'''transformer.blocks.{i}.mlp.fc1.weight''', f'''vilt.encoder.layer.{i}.intermediate.dense.weight''')			)
    rename_keys.append((f'''transformer.blocks.{i}.mlp.fc1.bias''', f'''vilt.encoder.layer.{i}.intermediate.dense.bias''')			)
    rename_keys.append((f'''transformer.blocks.{i}.mlp.fc2.weight''', f'''vilt.encoder.layer.{i}.output.dense.weight''')			)
    rename_keys.append((f'''transformer.blocks.{i}.mlp.fc2.bias''', f'''vilt.encoder.layer.{i}.output.dense.bias''')			)

  # embeddings
  rename_keys.extend(
      [
          # text embeddings
          ('text_embeddings.word_embeddings.weight', 'vilt.embeddings.text_embeddings.word_embeddings.weight'),
          (
              'text_embeddings.position_embeddings.weight',
              'vilt.embeddings.text_embeddings.position_embeddings.weight',
          ),
          ('text_embeddings.position_ids', 'vilt.embeddings.text_embeddings.position_ids'),
          (
              'text_embeddings.token_type_embeddings.weight',
              'vilt.embeddings.text_embeddings.token_type_embeddings.weight',
          ),
          ('text_embeddings.LayerNorm.weight', 'vilt.embeddings.text_embeddings.LayerNorm.weight'),
          ('text_embeddings.LayerNorm.bias', 'vilt.embeddings.text_embeddings.LayerNorm.bias'),
          # patch embeddings
          ('transformer.cls_token', 'vilt.embeddings.cls_token'),
          ('transformer.patch_embed.proj.weight', 'vilt.embeddings.patch_embeddings.projection.weight'),
          ('transformer.patch_embed.proj.bias', 'vilt.embeddings.patch_embeddings.projection.bias'),
          ('transformer.pos_embed', 'vilt.embeddings.position_embeddings'),
          # token type embeddings
          ('token_type_embeddings.weight', 'vilt.embeddings.token_type_embeddings.weight'),
      ]			)

  # final layernorm + pooler
  rename_keys.extend(
      [
          ('transformer.norm.weight', 'vilt.layernorm.weight'),
          ('transformer.norm.bias', 'vilt.layernorm.bias'),
          ('pooler.dense.weight', 'vilt.pooler.dense.weight'),
          ('pooler.dense.bias', 'vilt.pooler.dense.bias'),
      ]			)

  # classifier head(s)
  if vqa_model:
    # classification head
    rename_keys.extend(
        [
            ('vqa_classifier.0.weight', 'classifier.0.weight'),
            ('vqa_classifier.0.bias', 'classifier.0.bias'),
            ('vqa_classifier.1.weight', 'classifier.1.weight'),
            ('vqa_classifier.1.bias', 'classifier.1.bias'),
            ('vqa_classifier.3.weight', 'classifier.3.weight'),
            ('vqa_classifier.3.bias', 'classifier.3.bias'),
        ]			)
  elif nlvr_model:
    # classification head
    rename_keys.extend(
        [
            ('nlvr2_classifier.0.weight', 'classifier.0.weight'),
            ('nlvr2_classifier.0.bias', 'classifier.0.bias'),
            ('nlvr2_classifier.1.weight', 'classifier.1.weight'),
            ('nlvr2_classifier.1.bias', 'classifier.1.bias'),
            ('nlvr2_classifier.3.weight', 'classifier.3.weight'),
            ('nlvr2_classifier.3.bias', 'classifier.3.bias'),
        ]			)
  else:
    pass

  return rename_keys


def 				A	( lowercase   ,	lowercase			)       ->  Tuple:
  '''simple docstring'''

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


def 				A	( lowercase			)       ->  List[str]:
  '''simple docstring'''

  UpperCamelCase          =					['head.weight', 'head.bias']
  for k in ignore_keys:
    state_dict.pop(lowerCAmelCase_   ,	lowerCAmelCase_			)


def 				A	( lowercase   ,	lowercase   ,	lowercase			)       ->  Optional[int]:
  '''simple docstring'''

  UpperCamelCase          =					dct.pop(lowerCAmelCase_			)
  UpperCamelCase          =					val


@torch.no_grad()
def 				A	( lowercase   ,	lowercase			)       ->  Tuple:
  '''simple docstring'''

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

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

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

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

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

  # Verify outputs
  if mlm_model:
    UpperCamelCase          =					torch.Size([1, 11, 30_522]			)
    UpperCamelCase          =					torch.tensor([-1_2.5_0_6_1, -1_2.5_1_2_3, -1_2.5_1_7_4]			)
    assert outputs.logits.shape == expected_shape
    assert torch.allclose(outputs.logits[0, 0, :3]   ,	lowerCAmelCase_   ,	atol=1e-4			)

    # verify masked token prediction equals "cats"
    UpperCamelCase          =					outputs.logits[0, 4, :].argmax(-1			).item()
    assert tokenizer.decode([predicted_id]			) == "cats"
  elif vqa_model:
    UpperCamelCase          =					torch.Size([1, 3_129]			)
    UpperCamelCase          =					torch.tensor([-1_5.9_4_9_5, -1_8.1_4_7_2, -1_0.3_0_4_1]			)
    assert torch.allclose(outputs.logits[0, :3]   ,	lowerCAmelCase_   ,	atol=1e-4			)
    assert outputs.logits.shape == expected_shape
    assert torch.allclose(outputs.logits[0, 0, :3]   ,	lowerCAmelCase_   ,	atol=1e-4			)

    # verify vqa prediction equals "2"
    UpperCamelCase          =					outputs.logits.argmax(-1			).item()
    assert model.config.idalabel[predicted_idx] == "2"
  elif nlvr_model:
    UpperCamelCase          =					torch.Size([1, 2]			)
    UpperCamelCase          =					torch.tensor([-2.8_7_2_1, 2.1_2_9_1]			)
    assert torch.allclose(outputs.logits[0, :3]   ,	lowerCAmelCase_   ,	atol=1e-4			)
    assert outputs.logits.shape == expected_shape

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


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

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


import unittest

from transformers import GPTSwaTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_UpperCAmelCase :			Union[str, Any]              =	get_tests_dir("fixtures/test_sentencepiece_with_bytefallback.model")


@require_sentencepiece
@require_tokenizers
class     lowercase (		_SCREAMING_SNAKE_CASE  ,       unittest.TestCase       ):
		__lowercase	:	Optional[Any]     					=    GPTSwaTokenizer
		__lowercase	:	Optional[Any]     					=    False
		__lowercase	:	Union[str, Any]     					=    True
		__lowercase	:	Tuple     					=    False



		def __UpperCamelCase			(			self	)  -> Optional[int]:
				"""simple docstring"""
				super().setUp()

				# We have a SentencePiece fixture for testing
				UpperCamelCase          =					GPTSwaTokenizer(A_    ,				eos_token='<unk>'    ,				bos_token='<unk>'    ,				pad_token='<unk>'	)

				tokenizer.save_pretrained(self.tmpdirname	)



		def __UpperCamelCase			(			self    ,				A_	)  -> List[str]:
				"""simple docstring"""
				UpperCamelCase          =					'This is a test'
				UpperCamelCase          =					'This is a test'
				return input_text, output_text



		def __UpperCamelCase			(			self	)  -> int:
				"""simple docstring"""
				UpperCamelCase          =					'<s>'
				UpperCamelCase          =					1

				self.assertEqual(self.get_tokenizer()._convert_token_to_id(A_	)    ,				A_	)
				self.assertEqual(self.get_tokenizer()._convert_id_to_token(A_	)    ,				A_	)



		def __UpperCamelCase			(			self	)  -> Dict:
				"""simple docstring"""
				UpperCamelCase          =					list(self.get_tokenizer().get_vocab().keys()	)

				self.assertEqual(vocab_keys[0]    ,				'<unk>'	)
				self.assertEqual(vocab_keys[1]    ,				'<s>'	)
				self.assertEqual(vocab_keys[-1]    ,				'j'	)
				self.assertEqual(len(A_	)    ,				2_000	)



		def __UpperCamelCase			(			self	)  -> Any:
				"""simple docstring"""
				self.assertEqual(self.get_tokenizer().vocab_size    ,				2_000	)



		def __UpperCamelCase			(			self	)  -> int:
				"""simple docstring"""
				UpperCamelCase          =					GPTSwaTokenizer(A_	)

				UpperCamelCase          =					tokenizer.tokenize('This is a test'	)
				self.assertListEqual(A_    ,				['▁This', '▁is', '▁a', '▁t', 'est']	)

				self.assertListEqual(tokenizer.convert_tokens_to_ids(A_	)    ,				[465, 287, 265, 631, 842]	)

				UpperCamelCase          =					tokenizer.tokenize('I was born in 92000, and this is falsé.'	)
				# fmt: off
				self.assertListEqual(
				    A_    ,				['▁I', '▁was', '▁bor', 'n', '▁in', '▁', '<0x39>', '2', '0', '0', '0', ',', '▁and', '▁this', '▁is', '▁f', 'al', 's', '<0xC3>', '<0xA9>', '.']    ,				)
				# fmt: on

				UpperCamelCase          =					tokenizer.convert_tokens_to_ids(A_	)
				self.assertListEqual(
				    A_    ,				[262, 272, 1_525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260]    ,				)

				UpperCamelCase          =					tokenizer.convert_ids_to_tokens(A_	)
				# fmt: off
				self.assertListEqual(
				    A_    ,				['▁I', '▁was', '▁bor', 'n', '▁in', '▁', '<0x39>', '2', '0', '0', '0', ',', '▁and', '▁this', '▁is', '▁f', 'al', 's', '<0xC3>', '<0xA9>', '.']	)
				# fmt: on



		def __UpperCamelCase			(			self	)  -> Tuple:
				"""simple docstring"""
				UpperCamelCase          =					GPTSwaTokenizer(A_	)
				UpperCamelCase          =					['This is a test', 'I was born in 92000, and this is falsé.']
				UpperCamelCase          =					[
				    [465, 287, 265, 631, 842],
				    [262, 272, 1_525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260],
				]

				# Test that encode_fast returns the same as tokenize + convert_tokens_to_ids
				for text, expected_ids in zip(A_    ,				A_	):
						self.assertListEqual(tokenizer.encode_fast(A_	)    ,				A_	)

				# Test that decode_fast returns the input text
				for text, token_ids in zip(A_    ,				A_	):
						self.assertEqual(tokenizer.decode_fast(A_	)    ,				A_	)



		@slow
		def __UpperCamelCase			(			self	)  -> List[Any]:
				"""simple docstring"""
				UpperCamelCase          =					[
				    '<|python|>def fibonacci(n)\n    if n < 0:\n        print(\'Incorrect input\')',
				    'Hey there, how are you doing this fine day?',
				    'This is a text with a trailing spaces followed by a dot     .',
				    'Häj sväjs lillebrör! =)',
				    'Det är inget fel på Mr. Cool',
				]

				# fmt: off
				UpperCamelCase          =					{'input_ids': [[63_423, 5, 6_811, 14_954, 282, 816, 3_821, 63_466, 63_425, 63_462, 18, 63_978, 678, 301, 1_320, 63_423, 63_455, 63_458, 18, 63_982, 4_246, 3_940, 1_901, 47_789, 5_547, 18_994], [19_630, 1_100, 63_446, 1_342, 633, 544, 4_488, 593, 5_102, 2_416, 63_495, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1_652, 428, 268, 1_936, 515, 268, 58_593, 22_413, 9_106, 546, 268, 33_213, 63_979, 698, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [55_130, 63_450, 924, 63_449, 2_249, 4_062, 1_558, 318, 63_504, 21_498, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [509, 377, 2_827, 2_559, 332, 6_575, 63_443, 26_801, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], 'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}
				# fmt: on
				self.tokenizer_integration_test_util(
				    expected_encoding=A_    ,				model_name='AI-Sweden/gpt-sw3-126m'    ,				sequences=A_    ,				)
