





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 							__SCREAMING_SNAKE_CASE       (  _a ,		unittest.TestCase							):
							snake_case  :   Dict      =							CodeGenTokenizer
							snake_case  :   Dict      =							CodeGenTokenizerFast
							snake_case  :   Tuple      =							True
							snake_case  :   Optional[int]      =							{"""add_prefix_space""": True}
							snake_case  :   int      =							False






							def 					_lowerCamelCase		(    self	):
														super().setUp()

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

														UpperCamelCase__					  =      os.path.join(self.tmpdirname	,			VOCAB_FILES_NAMES["""vocab_file"""]	)
														UpperCamelCase__					  =      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(__lowerCAmelCase	) + """\n"""	)
														with open(self.merges_file	,			"""w"""	,			encoding="""utf-8"""	) as fp:
																					fp.write("""\n""".join(__lowerCAmelCase	)	)






							def 					_lowerCamelCase		(    self	,			**__lowerCAmelCase	):
														kwargs.update(self.special_tokens_map	)
														return CodeGenTokenizer.from_pretrained(self.tmpdirname	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			**__lowerCAmelCase	):
														kwargs.update(self.special_tokens_map	)
														return CodeGenTokenizerFast.from_pretrained(self.tmpdirname	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      """lower newer"""
														UpperCamelCase__					  =      """lower newer"""
														return input_text, output_text






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer(self.vocab_file	,			self.merges_file	,			**self.special_tokens_map	)
														UpperCamelCase__					  =      """lower newer"""
														UpperCamelCase__					  =      ["""\u0120low""", """er""", """\u0120""", """n""", """e""", """w""", """er"""]
														UpperCamelCase__					  =      tokenizer.tokenize(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														UpperCamelCase__					  =      tokens + [tokenizer.unk_token]
														UpperCamelCase__					  =      [14, 15, 10, 9, 3, 2, 15, 19]
														self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowerCAmelCase	)	,			__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	):
														if not self.test_rust_tokenizer:
																					return

														UpperCamelCase__					  =      self.get_tokenizer()
														UpperCamelCase__					  =      self.get_rust_tokenizer(add_prefix_space=__lowerCAmelCase	)

														UpperCamelCase__					  =      """lower newer"""

														# Testing tokenization
														UpperCamelCase__					  =      tokenizer.tokenize(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.tokenize(__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing conversion to ids without special tokens
														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	,			add_special_tokens=__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.encode(__lowerCAmelCase	,			add_special_tokens=__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing conversion to ids with special tokens
														UpperCamelCase__					  =      self.get_rust_tokenizer(add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	)
														UpperCamelCase__					  =      rust_tokenizer.encode(__lowerCAmelCase	)
														self.assertListEqual(__lowerCAmelCase	,			__lowerCAmelCase	)

														# Testing the unknown token
														UpperCamelCase__					  =      tokens + [rust_tokenizer.unk_token]
														UpperCamelCase__					  =      [14, 15, 10, 9, 3, 2, 15, 19]
														self.assertListEqual(rust_tokenizer.convert_tokens_to_ids(__lowerCAmelCase	)	,			__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			*__lowerCAmelCase	,			**__lowerCAmelCase	):
														# It's very difficult to mix/test pretokenization with byte-level
														# And get both CodeGen and Roberta to work at the same time (mostly an issue of adding a space before the string)
														pass






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase=15	):
														for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
																					with self.subTest(f"""{tokenizer.__class__.__name__} ({pretrained_name})"""	):
																												UpperCamelCase__					  =      self.rust_tokenizer_class.from_pretrained(__lowerCAmelCase	,			**__lowerCAmelCase	)

																												# Simple input
																												UpperCamelCase__					  =      """This is a simple input"""
																												UpperCamelCase__					  =      ["""This is a simple input 1""", """This is a simple input 2"""]
																												UpperCamelCase__					  =      ("""This is a simple input""", """This is a pair""")
																												UpperCamelCase__					  =      [
																												    ("""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(__lowerCAmelCase	,			tokenizer_r.encode	,			__lowerCAmelCase	,			max_length=__lowerCAmelCase	,			padding="""max_length"""	)

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

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

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

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

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






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained(self.tmpdirname	,			pad_token="""<pad>"""	)

														# Simple input
														UpperCamelCase__					  =      """This is a simple input"""
														UpperCamelCase__					  =      ["""This is a simple input looooooooong""", """This is a simple input"""]
														UpperCamelCase__					  =      ("""This is a simple input""", """This is a pair""")
														UpperCamelCase__					  =      [
														    ("""This is a simple input loooooong""", """This is a simple input"""),
														    ("""This is a simple pair loooooong""", """This is a simple pair"""),
														]

														UpperCamelCase__					  =      tokenizer.pad_token_id

														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding="""max_length"""	,			max_length=30	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding=__lowerCAmelCase	,			truncate=__lowerCAmelCase	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(*__lowerCAmelCase	,			padding="""max_length"""	,			max_length=60	,			return_tensors="""np"""	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	,			padding=__lowerCAmelCase	,			truncate=__lowerCAmelCase	,			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	):
														UpperCamelCase__					  =      """$$$"""
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained(self.tmpdirname	,			bos_token=__lowerCAmelCase	,			add_bos_token=__lowerCAmelCase	)

														UpperCamelCase__					  =      """This is a simple input"""
														UpperCamelCase__					  =      ["""This is a simple input 1""", """This is a simple input 2"""]

														UpperCamelCase__					  =      tokenizer.bos_token_id

														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	)
														UpperCamelCase__					  =      tokenizer(__lowerCAmelCase	)

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

														UpperCamelCase__					  =      tokenizer.decode(out_s.input_ids	)
														UpperCamelCase__					  =      tokenizer.batch_decode(out_sa.input_ids	)

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






							@slow
							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      CodeGenTokenizer.from_pretrained("""Salesforce/codegen-350M-mono"""	)

														UpperCamelCase__					  =      """\nif len_a > len_b:\n    result = a\nelse:\n    result = b\n\n\n\n#"""
														UpperCamelCase__					  =      """\nif len_a > len_b:      result = a\nelse:      result = b"""

														UpperCamelCase__					  =      tokenizer.encode(__lowerCAmelCase	)
														UpperCamelCase__					  =      ["""^#""", re.escape("""<|endoftext|>"""	), """^'''""", """^\"\"\"""", """\n\n\n"""]
														UpperCamelCase__					  =      tokenizer.decode(__lowerCAmelCase	,			truncate_before_pattern=__lowerCAmelCase	)
														self.assertEqual(__lowerCAmelCase	,			__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	):
														pass











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      _UpperCamelCase  (a__ :Tuple ,      a__ :List[str]=False							):






							"""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"""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"
														UpperCamelCase__					  =      [(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      _UpperCamelCase  (a__ :int ,      a__ :Tuple ,      a__ :Union[str, Any]=False							):






							"""simple docstring"""





							for i in range(config.num_hidden_layers							):
														if base_model:
																					UpperCamelCase__					  =      """"""
														else:
																					UpperCamelCase__					  =      """vit."""
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														UpperCamelCase__					  =      state_dict.pop(f"""module.blocks.{i}.attn.qkv.weight"""							)
														UpperCamelCase__					  =      state_dict.pop(f"""module.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      _UpperCamelCase  (a__ :int							):






							"""simple docstring"""





							UpperCamelCase__					  =      ["""head.weight""", """head.bias"""]
							for k in ignore_keys:
														state_dict.pop(a__ ,      a__							)





def      _UpperCamelCase  (a__ :Optional[int]							):






							"""simple docstring"""





							UpperCamelCase__					  =      [
							    """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(a__ ,      a__							)





def      _UpperCamelCase  (a__ :List[str] ,      a__ :Any ,      a__ :str							):






							"""simple docstring"""





							UpperCamelCase__					  =      dct.pop(a__							)
							UpperCamelCase__					  =      val





def      _UpperCamelCase  (a__ :Dict ,      a__ :Dict							):






							"""simple docstring"""





							UpperCamelCase__					  =      ViTMSNConfig()
							UpperCamelCase__					  =      1000

							UpperCamelCase__					  =      """datasets/huggingface/label-files"""
							UpperCamelCase__					  =      """imagenet-1k-id2label.json"""
							UpperCamelCase__					  =      json.load(open(hf_hub_download(a__ ,      a__							) ,      """r"""							)							)
							UpperCamelCase__					  =      {int(a__							): v for k, v in idalabel.items()}
							UpperCamelCase__					  =      idalabel
							UpperCamelCase__					  =      {v: k for k, v in idalabel.items()}

							if "s16" in checkpoint_url:
														UpperCamelCase__					  =      384
														UpperCamelCase__					  =      1536
														UpperCamelCase__					  =      6
							elif "l16" in checkpoint_url:
														UpperCamelCase__					  =      1024
														UpperCamelCase__					  =      4096
														UpperCamelCase__					  =      24
														UpperCamelCase__					  =      16
														UpperCamelCase__					  =      0.1
							elif "b4" in checkpoint_url:
														UpperCamelCase__					  =      4
							elif "l7" in checkpoint_url:
														UpperCamelCase__					  =      7
														UpperCamelCase__					  =      1024
														UpperCamelCase__					  =      4096
														UpperCamelCase__					  =      24
														UpperCamelCase__					  =      16
														UpperCamelCase__					  =      0.1

							UpperCamelCase__					  =      ViTMSNModel(a__							)

							UpperCamelCase__					  =      torch.hub.load_state_dict_from_url(a__ ,      map_location="""cpu"""							)["""target_encoder"""]

							UpperCamelCase__					  =      ViTImageProcessor(size=config.image_size							)

							remove_projection_head(a__							)
							UpperCamelCase__					  =      create_rename_keys(a__ ,      base_model=a__							)

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

							model.load_state_dict(a__							)
							model.eval()

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

							UpperCamelCase__					  =      Image.open(requests.get(a__ ,      stream=a__							).raw							)
							UpperCamelCase__					  =      ViTImageProcessor(
							    size=config.image_size ,      image_mean=a__ ,      image_std=a__							)
							UpperCamelCase__					  =      image_processor(images=a__ ,      return_tensors="""pt"""							)

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

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

							print(f"""Saving model to {pytorch_dump_folder_path}"""							)
							model.save_pretrained(a__							)

							print(f"""Saving image processor to {pytorch_dump_folder_path}"""							)
							image_processor.save_pretrained(a__							)


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)




