





import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_roberta import RobertaTokenizer


UpperCamelCase__           =			logging.get_logger(__name__)

UpperCamelCase__           =			{"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

UpperCamelCase__           =			{
    "vocab_file": {
        "roberta-base": "https://huggingface.co/roberta-base/resolve/main/vocab.json",
        "roberta-large": "https://huggingface.co/roberta-large/resolve/main/vocab.json",
        "roberta-large-mnli": "https://huggingface.co/roberta-large-mnli/resolve/main/vocab.json",
        "distilroberta-base": "https://huggingface.co/distilroberta-base/resolve/main/vocab.json",
        "roberta-base-openai-detector": "https://huggingface.co/roberta-base-openai-detector/resolve/main/vocab.json",
        "roberta-large-openai-detector": (
            "https://huggingface.co/roberta-large-openai-detector/resolve/main/vocab.json"
        ),
    },
    "merges_file": {
        "roberta-base": "https://huggingface.co/roberta-base/resolve/main/merges.txt",
        "roberta-large": "https://huggingface.co/roberta-large/resolve/main/merges.txt",
        "roberta-large-mnli": "https://huggingface.co/roberta-large-mnli/resolve/main/merges.txt",
        "distilroberta-base": "https://huggingface.co/distilroberta-base/resolve/main/merges.txt",
        "roberta-base-openai-detector": "https://huggingface.co/roberta-base-openai-detector/resolve/main/merges.txt",
        "roberta-large-openai-detector": (
            "https://huggingface.co/roberta-large-openai-detector/resolve/main/merges.txt"
        ),
    },
    "tokenizer_file": {
        "roberta-base": "https://huggingface.co/roberta-base/resolve/main/tokenizer.json",
        "roberta-large": "https://huggingface.co/roberta-large/resolve/main/tokenizer.json",
        "roberta-large-mnli": "https://huggingface.co/roberta-large-mnli/resolve/main/tokenizer.json",
        "distilroberta-base": "https://huggingface.co/distilroberta-base/resolve/main/tokenizer.json",
        "roberta-base-openai-detector": (
            "https://huggingface.co/roberta-base-openai-detector/resolve/main/tokenizer.json"
        ),
        "roberta-large-openai-detector": (
            "https://huggingface.co/roberta-large-openai-detector/resolve/main/tokenizer.json"
        ),
    },
}

UpperCamelCase__           =			{
    "roberta-base": 512,
    "roberta-large": 512,
    "roberta-large-mnli": 512,
    "distilroberta-base": 512,
    "roberta-base-openai-detector": 512,
    "roberta-large-openai-detector": 512,
}

class 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   Optional[Any]      =							VOCAB_FILES_NAMES
							snake_case  :   Optional[Any]      =							PRETRAINED_VOCAB_FILES_MAP
							snake_case  :   Any      =							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							snake_case  :   Dict      =							["""input_ids""", """attention_mask"""]
							snake_case  :   List[Any]      =							RobertaTokenizer






							def __init__(    self	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			__lowerCAmelCase=None	,			__lowerCAmelCase="replace"	,			__lowerCAmelCase="<s>"	,			__lowerCAmelCase="</s>"	,			__lowerCAmelCase="</s>"	,			__lowerCAmelCase="<s>"	,			__lowerCAmelCase="<unk>"	,			__lowerCAmelCase="<pad>"	,			__lowerCAmelCase="<mask>"	,			__lowerCAmelCase=False	,			__lowerCAmelCase=True	,			**__lowerCAmelCase	,			):
														super().__init__(
														    __lowerCAmelCase	,			__lowerCAmelCase	,			tokenizer_file=__lowerCAmelCase	,			errors=__lowerCAmelCase	,			bos_token=__lowerCAmelCase	,			eos_token=__lowerCAmelCase	,			sep_token=__lowerCAmelCase	,			cls_token=__lowerCAmelCase	,			unk_token=__lowerCAmelCase	,			pad_token=__lowerCAmelCase	,			mask_token=__lowerCAmelCase	,			add_prefix_space=__lowerCAmelCase	,			trim_offsets=__lowerCAmelCase	,			**__lowerCAmelCase	,			)

														UpperCamelCase__					  =      json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
														if pre_tok_state.get("""add_prefix_space"""	,			__lowerCAmelCase	) != add_prefix_space:
																					UpperCamelCase__					  =      getattr(__lowerCAmelCase	,			pre_tok_state.pop("""type"""	)	)
																					UpperCamelCase__					  =      add_prefix_space
																					UpperCamelCase__					  =      pre_tok_class(**__lowerCAmelCase	)

														UpperCamelCase__					  =      add_prefix_space

														UpperCamelCase__					  =      """post_processor"""
														UpperCamelCase__					  =      getattr(self.backend_tokenizer	,			__lowerCAmelCase	,			__lowerCAmelCase	)
														if tokenizer_component_instance:
																					UpperCamelCase__					  =      json.loads(tokenizer_component_instance.__getstate__()	)

																					# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
																					if "sep" in state:
																												UpperCamelCase__					  =      tuple(state["""sep"""]	)
																					if "cls" in state:
																												UpperCamelCase__					  =      tuple(state["""cls"""]	)

																					UpperCamelCase__					  =      False

																					if state.get("""add_prefix_space"""	,			__lowerCAmelCase	) != add_prefix_space:
																												UpperCamelCase__					  =      add_prefix_space
																												UpperCamelCase__					  =      True

																					if state.get("""trim_offsets"""	,			__lowerCAmelCase	) != trim_offsets:
																												UpperCamelCase__					  =      trim_offsets
																												UpperCamelCase__					  =      True

																					if changes_to_apply:
																												UpperCamelCase__					  =      getattr(__lowerCAmelCase	,			state.pop("""type"""	)	)
																												UpperCamelCase__					  =      component_class(**__lowerCAmelCase	)
																												setattr(self.backend_tokenizer	,			__lowerCAmelCase	,			__lowerCAmelCase	)






							@property
							def 					_lowerCamelCase		(    self	):
														if self._mask_token is None:
																					if self.verbose:
																												logger.error("""Using mask_token, but it is not set yet."""	)
																					return None
														return str(self._mask_token	)






							@mask_token.setter
							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      AddedToken(__lowerCAmelCase	,			lstrip=__lowerCAmelCase	,			rstrip=__lowerCAmelCase	) if isinstance(__lowerCAmelCase	,			__lowerCAmelCase	) else value
														UpperCamelCase__					  =      value






							def 					_lowerCamelCase		(    self	,			*__lowerCAmelCase	,			**__lowerCAmelCase	):
														UpperCamelCase__					  =      kwargs.get("""is_split_into_words"""	,			__lowerCAmelCase	)
														assert self.add_prefix_space or not is_split_into_words, (
														    f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
														    "to use it with pretokenized inputs."
														)

														return super()._batch_encode_plus(*__lowerCAmelCase	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			*__lowerCAmelCase	,			**__lowerCAmelCase	):
														UpperCamelCase__					  =      kwargs.get("""is_split_into_words"""	,			__lowerCAmelCase	)

														assert self.add_prefix_space or not is_split_into_words, (
														    f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
														    "to use it with pretokenized inputs."
														)

														return super()._encode_plus(*__lowerCAmelCase	,			**__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase = None	):
														UpperCamelCase__					  =      self._tokenizer.model.save(__lowerCAmelCase	,			name=__lowerCAmelCase	)
														return tuple(__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=None	):
														UpperCamelCase__					  =      [self.bos_token_id] + token_ids_a + [self.eos_token_id]
														if token_ids_a is None:
																					return output

														return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase = None	):
														UpperCamelCase__					  =      [self.sep_token_id]
														UpperCamelCase__					  =      [self.cls_token_id]

														if token_ids_a is None:
																					return len(cls + token_ids_a + sep	) * [0]
														return len(cls + token_ids_a + sep + sep + token_ids_a + sep	) * [0]











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)




