





'''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			(      __lowerCAmelCase       :							List[Any]      ,    __lowerCAmelCase       :							Optional[Any]=False		)		->  int:
   snake_case 						=   []
   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 						=   [(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			(      __lowerCAmelCase       :							str      ,    __lowerCAmelCase       :							Optional[int]      ,    __lowerCAmelCase       :							Any=False		)		->  Any:
   for i in range(config.num_hidden_layers		):
      if base_model:
         snake_case 						=   """"""
      else:
         snake_case 						=   """vit."""
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      snake_case 						=   state_dict.pop(F'''module.blocks.{i}.attn.qkv.weight'''		)
      snake_case 						=   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 						=   in_proj_weight[
          : config.hidden_size, :
      ]
      snake_case 						=   in_proj_bias[: config.hidden_size]
      snake_case 						=   in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      snake_case 						=   in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      snake_case 						=   in_proj_weight[
          -config.hidden_size :, :
      ]
      snake_case 						=   in_proj_bias[-config.hidden_size :]





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Dict		)		->  int:
   snake_case 						=   ["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(__lowerCAmelCase      ,    __lowerCAmelCase		)





def 	__lowerCamelCase			(      __lowerCAmelCase       :							List[str]		)		->  Optional[int]:
   # projection head is used in the self-supervised pre-training in MSN,
   # for downstream task it's not needed.
   snake_case 						=   [
       """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 	__lowerCamelCase			(      __lowerCAmelCase       :							Any      ,    __lowerCAmelCase       :							Dict      ,    __lowerCAmelCase       :							str		)		->  Dict:
   snake_case 						=   dct.pop(__lowerCAmelCase		)
   snake_case 						=   val





def 	__lowerCamelCase			(      __lowerCAmelCase       :							Optional[Any]      ,    __lowerCAmelCase       :							Optional[Any]		)		->  Dict:
   snake_case 						=   ViTMSNConfig()
   snake_case 						=   10_00

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

   if "s16" in checkpoint_url:
      snake_case 						=   3_84
      snake_case 						=   15_36
      snake_case 						=   6
   elif "l16" in checkpoint_url:
      snake_case 						=   10_24
      snake_case 						=   40_96
      snake_case 						=   24
      snake_case 						=   16
      snake_case 						=   0.1
   elif "b4" in checkpoint_url:
      snake_case 						=   4
   elif "l7" in checkpoint_url:
      snake_case 						=   7
      snake_case 						=   10_24
      snake_case 						=   40_96
      snake_case 						=   24
      snake_case 						=   16
      snake_case 						=   0.1

   snake_case 						=   ViTMSNModel(__lowerCAmelCase		)

   snake_case 						=   torch.hub.load_state_dict_from_url(__lowerCAmelCase      ,    map_location="""cpu"""		)["""target_encoder"""]

   snake_case 						=   ViTImageProcessor(size=config.image_size		)

   remove_projection_head(__lowerCAmelCase		)
   snake_case 						=   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()

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

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

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

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








'''simple docstring'''




import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
 import optax
 from flax.training.common_utils import onehot

 from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
 from transformers.models.ta.modeling_flax_ta import shift_tokens_right
@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			_lowerCAmelCase		(			unittest.TestCase	):






     """simple docstring"""






     @slow
     def      lowerCAmelCase			(   self							:			Union[str, Any]						)->	Optional[Any]:
        snake_case 						=   FlaxMTaForConditionalGeneration.from_pretrained("""google/mt5-small"""						)
        snake_case 						=   AutoTokenizer.from_pretrained("""google/mt5-small"""						)

        snake_case 						=   tokenizer("""Hello there"""   ,       return_tensors="""np"""						).input_ids
        snake_case 						=   tokenizer("""Hi I am"""   ,       return_tensors="""np"""						).input_ids

        snake_case 						=   shift_tokens_right(__snake_case   ,       model.config.pad_token_id   ,       model.config.decoder_start_token_id						)

        snake_case 						=   model(__snake_case   ,       decoder_input_ids=__snake_case						).logits
        snake_case 						=   optax.softmax_cross_entropy(__snake_case   ,       onehot(__snake_case   ,       logits.shape[-1]						)						).mean()

        snake_case 						=   -(labels.shape[-1] * loss.item())

        snake_case 						=   -84.91_27
        self.assertTrue(abs(mtf_score - EXPECTED_SCORE						) < 1e-4						)

