"""simple docstring"""



import json
import os
import re
import unittest

from transformers import CodeGenTokenizer, CodeGenTokenizerFast
from transformers.models.codegen.tokenization_codegen import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin
@require_tokenizers
class 						__UpperCAmelCase    (__A							,   unittest.TestCase     ):







			'''simple docstring'''
			_UpperCamelCase  :						str					     =				CodeGenTokenizer
			_UpperCamelCase  :						Union[str, Any]					     =				CodeGenTokenizerFast
			_UpperCamelCase  :						str					     =				True
			_UpperCamelCase  :						List[Any]					     =				{'add_prefix_space': True}
			_UpperCamelCase  :						List[Any]					     =				False





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										super().setUp()

										# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
										A__       :     Any								=			[
										    """l""",
										    """o""",
										    """w""",
										    """e""",
										    """r""",
										    """s""",
										    """t""",
										    """i""",
										    """d""",
										    """n""",
										    """\u0120""",
										    """\u0120l""",
										    """\u0120n""",
										    """\u0120lo""",
										    """\u0120low""",
										    """er""",
										    """\u0120lowest""",
										    """\u0120newer""",
										    """\u0120wider""",
										    """<unk>""",
										    """<|endoftext|>""",
										]
										A__       :     Tuple								=			dict(zip(snake_case_   ,      range(len(snake_case_							)							)							)							)
										A__       :     Dict								=			["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
										A__       :     Dict								=			{"""unk_token""": """<unk>"""}

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





			def 		lowerCamelCase					(			self   ,      **snake_case_							):







										'''simple docstring'''


										kwargs.update(self.special_tokens_map							)
										return CodeGenTokenizer.from_pretrained(self.tmpdirname   ,      **snake_case_							)





			def 		lowerCamelCase					(			self   ,      **snake_case_							):







										'''simple docstring'''


										kwargs.update(self.special_tokens_map							)
										return CodeGenTokenizerFast.from_pretrained(self.tmpdirname   ,      **snake_case_							)





			def 		lowerCamelCase					(			self   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Optional[int]								=			"""lower newer"""
										A__       :     int								=			"""lower newer"""
										return input_text, output_text





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[str]								=			CodeGenTokenizer(self.vocab_file   ,      self.merges_file   ,      **self.special_tokens_map							)
										A__       :     Optional[Any]								=			"""lower newer"""
										A__       :     Dict								=			["""\u0120low""", """er""", """\u0120""", """n""", """e""", """w""", """er"""]
										A__       :     List[str]								=			tokenizer.tokenize(snake_case_   ,      add_prefix_space=snake_case_							)
										self.assertListEqual(snake_case_   ,      snake_case_							)

										A__       :     Optional[Any]								=			tokens + [tokenizer.unk_token]
										A__       :     List[Any]								=			[14, 15, 10, 9, 3, 2, 15, 19]
										self.assertListEqual(tokenizer.convert_tokens_to_ids(snake_case_							)   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										if not self.test_rust_tokenizer:
																	return

										A__       :     Optional[int]								=			self.get_tokenizer()
										A__       :     str								=			self.get_rust_tokenizer(add_prefix_space=snake_case_							)

										A__       :     List[str]								=			"""lower newer"""

										# Testing tokenization
										A__       :     Dict								=			tokenizer.tokenize(snake_case_   ,      add_prefix_space=snake_case_							)
										A__       :     List[str]								=			rust_tokenizer.tokenize(snake_case_							)
										self.assertListEqual(snake_case_   ,      snake_case_							)

										# Testing conversion to ids without special tokens
										A__       :     Optional[int]								=			tokenizer.encode(snake_case_   ,      add_special_tokens=snake_case_   ,      add_prefix_space=snake_case_							)
										A__       :     Tuple								=			rust_tokenizer.encode(snake_case_   ,      add_special_tokens=snake_case_							)
										self.assertListEqual(snake_case_   ,      snake_case_							)

										# Testing conversion to ids with special tokens
										A__       :     Union[str, Any]								=			self.get_rust_tokenizer(add_prefix_space=snake_case_							)
										A__       :     Optional[Any]								=			tokenizer.encode(snake_case_   ,      add_prefix_space=snake_case_							)
										A__       :     Optional[int]								=			rust_tokenizer.encode(snake_case_							)
										self.assertListEqual(snake_case_   ,      snake_case_							)

										# Testing the unknown token
										A__       :     List[Any]								=			tokens + [rust_tokenizer.unk_token]
										A__       :     int								=			[14, 15, 10, 9, 3, 2, 15, 19]
										self.assertListEqual(rust_tokenizer.convert_tokens_to_ids(snake_case_							)   ,      snake_case_							)





			def 		lowerCamelCase					(			self   ,      *snake_case_   ,      **snake_case_							):







										'''simple docstring'''


										pass





			def 		lowerCamelCase					(			self   ,      snake_case_=15							):







										'''simple docstring'''


										for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
																	with self.subTest(F'''{tokenizer.__class__.__name__} ({pretrained_name})'''							):
																								A__       :     Union[str, Any]								=			self.rust_tokenizer_class.from_pretrained(snake_case_   ,      **snake_case_							)

																								# Simple input
																								A__       :     int								=			"""This is a simple input"""
																								A__       :     int								=			["""This is a simple input 1""", """This is a simple input 2"""]
																								A__       :     int								=			("""This is a simple input""", """This is a pair""")
																								A__       :     List[str]								=			[
																								    ("""This is a simple input 1""", """This is a simple input 2"""),
																								    ("""This is a simple pair 1""", """This is a simple pair 2"""),
																								]

																								# Simple input tests
																								self.assertRaises(snake_case_   ,      tokenizer_r.encode   ,      snake_case_   ,      max_length=snake_case_   ,      padding="""max_length"""							)

																								# Simple input
																								self.assertRaises(snake_case_   ,      tokenizer_r.encode_plus   ,      snake_case_   ,      max_length=snake_case_   ,      padding="""max_length"""							)

																								# Simple input
																								self.assertRaises(
																								    snake_case_   ,      tokenizer_r.batch_encode_plus   ,      snake_case_   ,      max_length=snake_case_   ,      padding="""max_length"""   ,      )

																								# Pair input
																								self.assertRaises(snake_case_   ,      tokenizer_r.encode   ,      snake_case_   ,      max_length=snake_case_   ,      padding="""max_length"""							)

																								# Pair input
																								self.assertRaises(snake_case_   ,      tokenizer_r.encode_plus   ,      snake_case_   ,      max_length=snake_case_   ,      padding="""max_length"""							)

																								# Pair input
																								self.assertRaises(
																								    snake_case_   ,      tokenizer_r.batch_encode_plus   ,      snake_case_   ,      max_length=snake_case_   ,      padding="""max_length"""   ,      )





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[str]								=			CodeGenTokenizer.from_pretrained(self.tmpdirname   ,      pad_token="""<pad>"""							)

										# Simple input
										A__       :     Dict								=			"""This is a simple input"""
										A__       :     Optional[int]								=			["""This is a simple input looooooooong""", """This is a simple input"""]
										A__       :     Dict								=			("""This is a simple input""", """This is a pair""")
										A__       :     Tuple								=			[
										    ("""This is a simple input loooooong""", """This is a simple input"""),
										    ("""This is a simple pair loooooong""", """This is a simple pair"""),
										]

										A__       :     List[str]								=			tokenizer.pad_token_id

										A__       :     List[str]								=			tokenizer(snake_case_   ,      padding="""max_length"""   ,      max_length=30   ,      return_tensors="""np"""							)
										A__       :     str								=			tokenizer(snake_case_   ,      padding=snake_case_   ,      truncate=snake_case_   ,      return_tensors="""np"""							)
										A__       :     Optional[int]								=			tokenizer(*snake_case_   ,      padding="""max_length"""   ,      max_length=60   ,      return_tensors="""np"""							)
										A__       :     Optional[int]								=			tokenizer(snake_case_   ,      padding=snake_case_   ,      truncate=snake_case_   ,      return_tensors="""np"""							)

										# s
										# test single string max_length padding
										self.assertEqual(out_s["""input_ids"""].shape[-1]   ,      30							)
										self.assertTrue(pad_token_id in out_s["""input_ids"""]							)
										self.assertTrue(0 in out_s["""attention_mask"""]							)

										# s2
										# test automatic padding
										self.assertEqual(out_sa["""input_ids"""].shape[-1]   ,      33							)
										# long slice doesn't have padding
										self.assertFalse(pad_token_id in out_sa["""input_ids"""][0]							)
										self.assertFalse(0 in out_sa["""attention_mask"""][0]							)
										# short slice does have padding
										self.assertTrue(pad_token_id in out_sa["""input_ids"""][1]							)
										self.assertTrue(0 in out_sa["""attention_mask"""][1]							)

										# p
										# test single pair max_length padding
										self.assertEqual(out_p["""input_ids"""].shape[-1]   ,      60							)
										self.assertTrue(pad_token_id in out_p["""input_ids"""]							)
										self.assertTrue(0 in out_p["""attention_mask"""]							)

										# p2
										# test automatic padding pair
										self.assertEqual(out_pa["""input_ids"""].shape[-1]   ,      52							)
										# long slice pair doesn't have padding
										self.assertFalse(pad_token_id in out_pa["""input_ids"""][0]							)
										self.assertFalse(0 in out_pa["""attention_mask"""][0]							)
										# short slice pair does have padding
										self.assertTrue(pad_token_id in out_pa["""input_ids"""][1]							)
										self.assertTrue(0 in out_pa["""attention_mask"""][1]							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     Optional[int]								=			"""$$$"""
										A__       :     List[str]								=			CodeGenTokenizer.from_pretrained(self.tmpdirname   ,      bos_token=snake_case_   ,      add_bos_token=snake_case_							)

										A__       :     Union[str, Any]								=			"""This is a simple input"""
										A__       :     Optional[int]								=			["""This is a simple input 1""", """This is a simple input 2"""]

										A__       :     List[Any]								=			tokenizer.bos_token_id

										A__       :     int								=			tokenizer(snake_case_							)
										A__       :     Union[str, Any]								=			tokenizer(snake_case_							)

										self.assertEqual(out_s.input_ids[0]   ,      snake_case_							)
										self.assertTrue(all(o[0] == bos_token_id for o in out_sa.input_ids							)							)

										A__       :     Dict								=			tokenizer.decode(out_s.input_ids							)
										A__       :     Optional[Any]								=			tokenizer.batch_decode(out_sa.input_ids							)

										self.assertEqual(decode_s.split()[0]   ,      snake_case_							)
										self.assertTrue(all(d.split()[0] == bos_token for d in decode_sa							)							)





			@slow
			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										A__       :     List[Any]								=			CodeGenTokenizer.from_pretrained("""Salesforce/codegen-350M-mono"""							)

										A__       :     List[str]								=			"""\nif len_a > len_b:\n    result = a\nelse:\n    result = b\n\n\n\n#"""
										A__       :     List[str]								=			"""\nif len_a > len_b:      result = a\nelse:      result = b"""

										A__       :     List[Any]								=			tokenizer.encode(snake_case_							)
										A__       :     Tuple								=			["""^#""", re.escape("""<|endoftext|>"""							), """^'''""", """^\"\"\"""", """\n\n\n"""]
										A__       :     Union[str, Any]								=			tokenizer.decode(snake_case_   ,      truncate_before_pattern=snake_case_							)
										self.assertEqual(snake_case_   ,      snake_case_							)





			def 		lowerCamelCase					(			self							):







										'''simple docstring'''


										pass






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

def  _A(		lowerCAmelCase							):
							A__       :     str								=			["""head.weight""", """head.bias"""]
							for k in ignore_keys:
														state_dict.pop(lowerCAmelCase			,  lowerCAmelCase							)

def  _A(		lowerCAmelCase							):
							# projection head is used in the self-supervised pre-training in MSN,
							# for downstream task it's not needed.
							A__       :     Dict								=			[
							    """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(lowerCAmelCase			,  lowerCAmelCase							)

def  _A(		lowerCAmelCase			,  lowerCAmelCase			,  lowerCAmelCase							):
							A__       :     Optional[Any]								=			dct.pop(lowerCAmelCase							)
							A__       :     List[Any]								=			val

def  _A(		lowerCAmelCase			,  lowerCAmelCase							):
							A__       :     str								=			ViTMSNConfig()
							A__       :     List[str]								=			1000

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

							if "s16" in checkpoint_url:
														A__       :     Tuple								=			384
														A__       :     int								=			1536
														A__       :     Union[str, Any]								=			6
							elif "l16" in checkpoint_url:
														A__       :     int								=			1024
														A__       :     int								=			4096
														A__       :     Any								=			24
														A__       :     int								=			16
														A__       :     Union[str, Any]								=			0.1
							elif "b4" in checkpoint_url:
														A__       :     Dict								=			4
							elif "l7" in checkpoint_url:
														A__       :     List[str]								=			7
														A__       :     Optional[int]								=			1024
														A__       :     Optional[int]								=			4096
														A__       :     Any								=			24
														A__       :     Union[str, Any]								=			16
														A__       :     Tuple								=			0.1

							A__       :     List[str]								=			ViTMSNModel(lowerCAmelCase							)

							A__       :     Any								=			torch.hub.load_state_dict_from_url(lowerCAmelCase			,  map_location="""cpu"""							)["""target_encoder"""]

							A__       :     List[str]								=			ViTImageProcessor(size=config.image_size							)

							remove_projection_head(lowerCAmelCase							)
							A__       :     List[str]								=			create_rename_keys(lowerCAmelCase			,  base_model=lowerCAmelCase							)

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

							model.load_state_dict(lowerCAmelCase							)
							model.eval()

							A__       :     str								=			"""http://images.cocodataset.org/val2017/000000039769.jpg"""

							A__       :     int								=			Image.open(requests.get(lowerCAmelCase			,  stream=lowerCAmelCase							).raw							)
							A__       :     List[Any]								=			ViTImageProcessor(
							    size=config.image_size			,  image_mean=lowerCAmelCase			,  image_std=lowerCAmelCase							)
							A__       :     Any								=			image_processor(images=lowerCAmelCase			,  return_tensors="""pt"""							)

							# forward pass
							torch.manual_seed(2							)
							A__       :     Tuple								=			model(**lowerCAmelCase							)
							A__       :     Union[str, Any]								=			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:
														A__       :     Union[str, Any]								=			torch.tensor([[-1.0915, -1.4876, -1.1809]]							)
							elif "b16" in checkpoint_url:
														A__       :     int								=			torch.tensor([[14.2889, -18.9045, 11.7281]]							)
							elif "l16" in checkpoint_url:
														A__       :     List[Any]								=			torch.tensor([[41.5028, -22.8681, 45.6475]]							)
							elif "b4" in checkpoint_url:
														A__       :     Optional[Any]								=			torch.tensor([[-4.3868, 5.2932, -0.4137]]							)
							else:
														A__       :     str								=			torch.tensor([[-0.1792, -0.6465, 2.4263]]							)

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

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

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


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

				_UpperCamelCase								=		parser.parse_args()
				convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)





