



"""simple docstring"""





import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			int=False    ):


						'''simple docstring'''






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

						# projection layer + position embeddings
						rename_keys.extend(
						    [
						        ("""module.cls_token""", """vit.embeddings.cls_token"""),
						        ("""module.patch_embed.proj.weight""", """vit.embeddings.patch_embeddings.projection.weight"""),
						        ("""module.patch_embed.proj.bias""", """vit.embeddings.patch_embeddings.projection.bias"""),
						        ("""module.pos_embed""", """vit.embeddings.position_embeddings"""),
						    ]    )

						if base_model:
												# layernorm + pooler
												rename_keys.extend(
												    [
												        ("""module.norm.weight""", """layernorm.weight"""),
												        ("""module.norm.bias""", """layernorm.bias"""),
												    ]    )

												# if just the base model, we should remove "vit" from all keys that start with "vit"
												snake_case_   : int						=			[(pair[0], pair[1][4:]) if pair[1].startswith("""vit"""    ) else pair for pair in rename_keys]
						else:
												# layernorm + classification head
												rename_keys.extend(
												    [
												        ("""norm.weight""", """vit.layernorm.weight"""),
												        ("""norm.bias""", """vit.layernorm.bias"""),
												        ("""head.weight""", """classifier.weight"""),
												        ("""head.bias""", """classifier.bias"""),
												    ]    )

						return rename_keys
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			str					, __UpperCamelCase      :			Tuple=False    ):


						'''simple docstring'''






						for i in range(config.num_hidden_layers    ):
												if base_model:
																		snake_case_   : Any						=			""""""
												else:
																		snake_case_   : List[Any]						=			"""vit."""
												# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
												snake_case_   : int						=			state_dict.pop(F'module.blocks.{i}.attn.qkv.weight'    )
												snake_case_   : int						=			state_dict.pop(F'module.blocks.{i}.attn.qkv.bias'    )
												# next, add query, keys and values (in that order) to the state dict
												snake_case_   : List[str]						=			in_proj_weight[
												    : config.hidden_size, :
												]
												snake_case_   : List[str]						=			in_proj_bias[: config.hidden_size]
												snake_case_   : Any						=			in_proj_weight[
												    config.hidden_size : config.hidden_size * 2, :
												]
												snake_case_   : List[str]						=			in_proj_bias[
												    config.hidden_size : config.hidden_size * 2
												]
												snake_case_   : int						=			in_proj_weight[
												    -config.hidden_size :, :
												]
												snake_case_   : Union[str, Any]						=			in_proj_bias[-config.hidden_size :]
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[Any]    ):


						'''simple docstring'''






						snake_case_   : int						=			["""head.weight""", """head.bias"""]
						for k in ignore_keys:
												state_dict.pop(__UpperCamelCase					, __UpperCamelCase    )
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Tuple    ):


						'''simple docstring'''






						snake_case_   : List[Any]						=			[
						    """module.fc.fc1.weight""",
						    """module.fc.fc1.bias""",
						    """module.fc.bn1.weight""",
						    """module.fc.bn1.bias""",
						    """module.fc.bn1.running_mean""",
						    """module.fc.bn1.running_var""",
						    """module.fc.bn1.num_batches_tracked""",
						    """module.fc.fc2.weight""",
						    """module.fc.fc2.bias""",
						    """module.fc.bn2.weight""",
						    """module.fc.bn2.bias""",
						    """module.fc.bn2.running_mean""",
						    """module.fc.bn2.running_var""",
						    """module.fc.bn2.num_batches_tracked""",
						    """module.fc.fc3.weight""",
						    """module.fc.fc3.bias""",
						]
						for k in ignore_keys:
												state_dict.pop(__UpperCamelCase					, __UpperCamelCase    )
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Union[str, Any]					, __UpperCamelCase      :			int					, __UpperCamelCase      :			List[Any]    ):


						'''simple docstring'''






						snake_case_   : Union[str, Any]						=			dct.pop(__UpperCamelCase    )
						snake_case_   : List[Any]						=			val
def 	__lowerCAmelCase						(		__UpperCamelCase      :			int					, __UpperCamelCase      :			Tuple    ):


						'''simple docstring'''






						snake_case_   : Tuple						=			ViTMSNConfig()
						snake_case_   : Optional[Any]						=			1_0_0_0

						snake_case_   : Dict						=			"""datasets/huggingface/label-files"""
						snake_case_   : Optional[int]						=			"""imagenet-1k-id2label.json"""
						snake_case_   : int						=			json.load(open(hf_hub_download(__UpperCamelCase					, __UpperCamelCase    )					, """r"""    )    )
						snake_case_   : Dict						=			{int(__UpperCamelCase    ): v for k, v in idalabel.items()}
						snake_case_   : List[Any]						=			idalabel
						snake_case_   : Dict						=			{v: k for k, v in idalabel.items()}

						if "s16" in checkpoint_url:
												snake_case_   : Dict						=			3_8_4
												snake_case_   : str						=			1_5_3_6
												snake_case_   : Tuple						=			6
						elif "l16" in checkpoint_url:
												snake_case_   : int						=			1_0_2_4
												snake_case_   : Tuple						=			4_0_9_6
												snake_case_   : Union[str, Any]						=			2_4
												snake_case_   : Union[str, Any]						=			1_6
												snake_case_   : Optional[Any]						=			0.1
						elif "b4" in checkpoint_url:
												snake_case_   : str						=			4
						elif "l7" in checkpoint_url:
												snake_case_   : List[Any]						=			7
												snake_case_   : List[str]						=			1_0_2_4
												snake_case_   : List[str]						=			4_0_9_6
												snake_case_   : Optional[Any]						=			2_4
												snake_case_   : List[Any]						=			1_6
												snake_case_   : int						=			0.1

						snake_case_   : List[Any]						=			ViTMSNModel(__UpperCamelCase    )

						snake_case_   : str						=			torch.hub.load_state_dict_from_url(__UpperCamelCase					, map_location="""cpu"""    )["""target_encoder"""]

						snake_case_   : Optional[int]						=			ViTImageProcessor(size=config.image_size    )

						remove_projection_head(__UpperCamelCase    )
						snake_case_   : Optional[Any]						=			create_rename_keys(__UpperCamelCase					, base_model=__UpperCamelCase    )

						for src, dest in rename_keys:
												rename_key(__UpperCamelCase					, __UpperCamelCase					, __UpperCamelCase    )
						read_in_q_k_v(__UpperCamelCase					, __UpperCamelCase					, base_model=__UpperCamelCase    )

						model.load_state_dict(__UpperCamelCase    )
						model.eval()

						snake_case_   : Union[str, Any]						=			"""http://images.cocodataset.org/val2017/000000039769.jpg"""

						snake_case_   : List[str]						=			Image.open(requests.get(__UpperCamelCase					, stream=__UpperCamelCase    ).raw    )
						snake_case_   : Optional[Any]						=			ViTImageProcessor(
						    size=config.image_size					, image_mean=__UpperCamelCase					, image_std=__UpperCamelCase    )
						snake_case_   : List[Any]						=			image_processor(images=__UpperCamelCase					, return_tensors="""pt"""    )

						# forward pass
						torch.manual_seed(2    )
						snake_case_   : Union[str, Any]						=			model(**__UpperCamelCase    )
						snake_case_   : int						=			outputs.last_hidden_state

						# The following Colab Notebook was used to generate these outputs:
						# https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
						if "s16" in checkpoint_url:
												snake_case_   : Any						=			torch.tensor([[-1.0_915, -1.4_876, -1.1_809]]    )
						elif "b16" in checkpoint_url:
												snake_case_   : List[Any]						=			torch.tensor([[14.2_889, -18.9_045, 11.7_281]]    )
						elif "l16" in checkpoint_url:
												snake_case_   : List[Any]						=			torch.tensor([[41.5_028, -22.8_681, 45.6_475]]    )
						elif "b4" in checkpoint_url:
												snake_case_   : Tuple						=			torch.tensor([[-4.3_868, 5.2_932, -0.4_137]]    )
						else:
												snake_case_   : Union[str, Any]						=			torch.tensor([[-0.1_792, -0.6_465, 2.4_263]]    )

						# verify logits
						assert torch.allclose(last_hidden_state[:, 0, :3]					, __UpperCamelCase					, atol=1E-4    )

						print(F'Saving model to {pytorch_dump_folder_path}'    )
						model.save_pretrained(__UpperCamelCase    )

						print(F'Saving image processor to {pytorch_dump_folder_path}'    )
						image_processor.save_pretrained(__UpperCamelCase    )


if __name__ == "__main__":
					__lowerCAmelCase       :					int      					=					argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--checkpoint_url''',
					    default='''https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar''',
					    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.'''
					)

					__lowerCAmelCase       :					Any      					=					parser.parse_args()
					convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






"""simple docstring"""





import json
import os
import unittest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin



@require_tokenizers
class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       MgpstrTokenizer
		_lowerCamelCase			  =       False
		_lowerCamelCase			  =       {}
		_lowerCamelCase			  =       False





		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								super().setUp()

								# fmt: off
								snake_case_   : Optional[Any]						=			["""[GO]""", """[s]""", """0""", """1""", """2""", """3""", """4""", """5""", """6""", """7""", """8""", """9""", """a""", """b""", """c""", """d""", """e""", """f""", """g""", """h""", """i""", """j""", """k""", """l""", """m""", """n""", """o""", """p""", """q""", """r""", """s""", """t""", """u""", """v""", """w""", """x""", """y""", """z"""]
								# fmt: on
								snake_case_   : str						=			dict(zip(_lowercase							,    range(len(_lowercase   )   )   )   )

								snake_case_   : Tuple						=			os.path.join(self.tmpdirname							,    VOCAB_FILES_NAMES["""vocab_file"""]   )
								with open(self.vocab_file							,    """w"""							,    encoding="""utf-8"""   ) as fp:
														fp.write(json.dumps(_lowercase   ) + """\n"""   )





		def    UpperCAmelCase__							(					self							,    **_lowercase   ) ->						Union[str, Any]:







								'''simple docstring'''

								return MgpstrTokenizer.from_pretrained(self.tmpdirname							,    **_lowercase   )





		def    UpperCAmelCase__							(					self							,    _lowercase   ) ->						Optional[Any]:







								'''simple docstring'''

								snake_case_   : Tuple						=			"""tester"""
								snake_case_   : Tuple						=			"""tester"""
								return input_text, output_text





		@unittest.skip("""MGP-STR always lower cases letters."""   )
		def    UpperCAmelCase__							(					self   ) ->						List[Any]:







								'''simple docstring'''

								pass





		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : Dict						=			self.get_tokenizers(do_lower_case=_lowercase   )
								for tokenizer in tokenizers:
														with self.subTest(f'{tokenizer.__class__.__name__}'   ):
																				snake_case_   : Optional[Any]						=			"""[SPECIAL_TOKEN]"""

																				tokenizer.add_special_tokens({"""cls_token""": special_token}   )
																				snake_case_   : Union[str, Any]						=			tokenizer.encode([special_token]							,    add_special_tokens=_lowercase   )
																				self.assertEqual(len(_lowercase   )							,    1   )

																				snake_case_   : List[Any]						=			tokenizer.decode(_lowercase							,    skip_special_tokens=_lowercase   )
																				self.assertTrue(special_token not in decoded   )





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : List[Any]						=			self.get_tokenizers()
								for tokenizer in tokenizers:
														with self.subTest(f'{tokenizer.__class__.__name__}'   ):
																				snake_case_			,		snake_case_   : Union[str, Any]						=			self.get_input_output_texts(_lowercase   )

																				snake_case_   : List[Any]						=			tokenizer.tokenize(_lowercase   )
																				snake_case_   : Union[str, Any]						=			tokenizer.convert_tokens_to_ids(_lowercase   )
																				snake_case_   : Union[str, Any]						=			tokenizer.encode(_lowercase							,    add_special_tokens=_lowercase   )
																				self.assertListEqual(_lowercase							,    _lowercase   )

																				snake_case_   : List[str]						=			tokenizer.convert_ids_to_tokens(_lowercase   )
																				self.assertNotEqual(len(_lowercase   )							,    0   )
																				snake_case_   : str						=			tokenizer.decode(_lowercase   )
																				self.assertIsInstance(_lowercase							,    _lowercase   )

																				self.assertEqual(text_a.replace(""" """							,    """"""   )							,    _lowercase   )





		@unittest.skip("""MGP-STR tokenizer only handles one sequence."""   )
		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								pass





		@unittest.skip("""inputs cannot be pretokenized in MgpstrTokenizer"""   )
		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								pass

