


'''simple docstring'''



import shutil
import tempfile
import unittest

from transformers import (
    SPIECE_UNDERLINE,
    AddedToken,
    BatchEncoding,
    NllbTokenizer,
    NllbTokenizerFast,
    is_torch_available,
)
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


A__		:					Dict												=  get_tests_dir("""fixtures/test_sentencepiece.model""")


if is_torch_available():
      from transformers.models.mam_aaa.modeling_mam_aaa import shift_tokens_right

A__		:					int												=  256047
A__		:					List[str]												=  256145





@require_sentencepiece
@require_tokenizers
class    UpperCAmelCase_  (_SCREAMING_SNAKE_CASE					,		unittest.TestCase  ):





    """simple docstring"""






    lowerCamelCase  :					Tuple							     =			NllbTokenizer
    lowerCamelCase  :					int							     =			NllbTokenizerFast
    lowerCamelCase  :					str							     =			True
    lowerCamelCase  :					Dict							     =			True
    lowerCamelCase  :					Union[str, Any]							     =			{}







    def       lowercase_				(							self	)				->     Tuple:
         super().setUp()

         # We have a SentencePiece fixture for testing
         __lowerCamelCase   :	Optional[Any]								=		NllbTokenizer(_lowerCAmelCase			,     keep_accents=_lowerCAmelCase	)
         tokenizer.save_pretrained(self.tmpdirname	)







    def       lowercase_				(							self	)				->     Optional[int]:
         __lowerCamelCase   :	List[str]								=		NllbTokenizer(_lowerCAmelCase			,     keep_accents=_lowerCAmelCase	)

         __lowerCamelCase   :	List[str]								=		tokenizer.tokenize('This is a test'	)
         self.assertListEqual(_lowerCAmelCase			,     ['▁This', '▁is', '▁a', '▁t', 'est']	)

         self.assertListEqual(
             tokenizer.convert_tokens_to_ids(_lowerCAmelCase	)			,     [value + tokenizer.fairseq_offset for value in [2_85, 46, 10, 1_70, 3_82]]			,     )

         __lowerCamelCase   :	List[Any]								=		tokenizer.tokenize('I was born in 92000, and this is falsé.'	)
         self.assertListEqual(
             _lowerCAmelCase			,     [
                 SPIECE_UNDERLINE + 'I',
                 SPIECE_UNDERLINE + 'was',
                 SPIECE_UNDERLINE + 'b',
                 'or',
                 'n',
                 SPIECE_UNDERLINE + 'in',
                 SPIECE_UNDERLINE + '',
                 '9',
                 '2',
                 '0',
                 '0',
                 '0',
                 ',',
                 SPIECE_UNDERLINE + 'and',
                 SPIECE_UNDERLINE + 'this',
                 SPIECE_UNDERLINE + 'is',
                 SPIECE_UNDERLINE + 'f',
                 'al',
                 's',
                 'é',
                 '.',
             ]			,     )
         __lowerCamelCase   :	Dict								=		tokenizer.convert_tokens_to_ids(_lowerCAmelCase	)
         self.assertListEqual(
             _lowerCAmelCase			,     [
                 value + tokenizer.fairseq_offset
                 for value in [8, 21, 84, 55, 24, 19, 7, 2, 6_02, 3_47, 3_47, 3_47, 3, 12, 66, 46, 72, 80, 6, 2, 4]
             ]			,     )

         __lowerCamelCase   :	str								=		tokenizer.convert_ids_to_tokens(_lowerCAmelCase	)
         self.assertListEqual(
             _lowerCAmelCase			,     [
                 SPIECE_UNDERLINE + 'I',
                 SPIECE_UNDERLINE + 'was',
                 SPIECE_UNDERLINE + 'b',
                 'or',
                 'n',
                 SPIECE_UNDERLINE + 'in',
                 SPIECE_UNDERLINE + '',
                 '<unk>',
                 '2',
                 '0',
                 '0',
                 '0',
                 ',',
                 SPIECE_UNDERLINE + 'and',
                 SPIECE_UNDERLINE + 'this',
                 SPIECE_UNDERLINE + 'is',
                 SPIECE_UNDERLINE + 'f',
                 'al',
                 's',
                 '<unk>',
                 '.',
             ]			,     )







    def       lowercase_				(							self	)				->     List[Any]:
         __lowerCamelCase   :	List[str]								=		(self.rust_tokenizer_class, 'hf-internal-testing/tiny-random-nllb', {})
         for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
              with self.subTest(f'{tokenizer.__class__.__name__} ({pretrained_name})'	):
                   __lowerCamelCase   :	str								=		self.rust_tokenizer_class.from_pretrained(_lowerCAmelCase			,     **_lowerCAmelCase	)
                   __lowerCamelCase   :	int								=		self.tokenizer_class.from_pretrained(_lowerCAmelCase			,     **_lowerCAmelCase	)

                   __lowerCamelCase   :	List[Any]								=		tempfile.mkdtemp()

                   __lowerCamelCase   :	Union[str, Any]								=		tokenizer_r.save_pretrained(_lowerCAmelCase	)
                   __lowerCamelCase   :	Tuple								=		tokenizer_p.save_pretrained(_lowerCAmelCase	)

                   # Checks it save with the same files + the tokenizer.json file for the fast one
                   self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files	)	)
                   __lowerCamelCase   :	Tuple								=		tuple(f for f in tokenizer_r_files if 'tokenizer.json' not in f	)
                   self.assertSequenceEqual(_lowerCAmelCase			,     _lowerCAmelCase	)

                   # Checks everything loads correctly in the same way
                   __lowerCamelCase   :	List[Any]								=		tokenizer_r.from_pretrained(_lowerCAmelCase	)
                   __lowerCamelCase   :	Any								=		tokenizer_p.from_pretrained(_lowerCAmelCase	)

                   # Check special tokens are set accordingly on Rust and Python
                   for key in tokenizer_pp.special_tokens_map:
                        self.assertTrue(hasattr(_lowerCAmelCase			,     _lowerCAmelCase	)	)

                   shutil.rmtree(_lowerCAmelCase	)

                   # Save tokenizer rust, legacy_format=True
                   __lowerCamelCase   :	Dict								=		tempfile.mkdtemp()

                   __lowerCamelCase   :	int								=		tokenizer_r.save_pretrained(_lowerCAmelCase			,     legacy_format=_lowerCAmelCase	)
                   __lowerCamelCase   :	int								=		tokenizer_p.save_pretrained(_lowerCAmelCase	)

                   # Checks it save with the same files
                   self.assertSequenceEqual(_lowerCAmelCase			,     _lowerCAmelCase	)

                   # Checks everything loads correctly in the same way
                   __lowerCamelCase   :	List[str]								=		tokenizer_r.from_pretrained(_lowerCAmelCase	)
                   __lowerCamelCase   :	Any								=		tokenizer_p.from_pretrained(_lowerCAmelCase	)

                   # Check special tokens are set accordingly on Rust and Python
                   for key in tokenizer_pp.special_tokens_map:
                        self.assertTrue(hasattr(_lowerCAmelCase			,     _lowerCAmelCase	)	)

                   shutil.rmtree(_lowerCAmelCase	)

                   # Save tokenizer rust, legacy_format=False
                   __lowerCamelCase   :	Tuple								=		tempfile.mkdtemp()

                   __lowerCamelCase   :	Dict								=		tokenizer_r.save_pretrained(_lowerCAmelCase			,     legacy_format=_lowerCAmelCase	)
                   __lowerCamelCase   :	Union[str, Any]								=		tokenizer_p.save_pretrained(_lowerCAmelCase	)

                   # Checks it saved the tokenizer.json file
                   self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files	)	)

                   # Checks everything loads correctly in the same way
                   __lowerCamelCase   :	Tuple								=		tokenizer_r.from_pretrained(_lowerCAmelCase	)
                   __lowerCamelCase   :	Union[str, Any]								=		tokenizer_p.from_pretrained(_lowerCAmelCase	)

                   # Check special tokens are set accordingly on Rust and Python
                   for key in tokenizer_pp.special_tokens_map:
                        self.assertTrue(hasattr(_lowerCAmelCase			,     _lowerCAmelCase	)	)

                   shutil.rmtree(_lowerCAmelCase	)







    @require_torch
    def       lowercase_				(							self	)				->     List[Any]:
         if not self.test_seqaseq:
              return

         __lowerCamelCase   :	List[Any]								=		self.get_tokenizers()
         for tokenizer in tokenizers:
              with self.subTest(f'{tokenizer.__class__.__name__}'	):
                   # Longer text that will definitely require truncation.
                   __lowerCamelCase   :	Union[str, Any]								=		[
                       ' UN Chief Says There Is No Military Solution in Syria',
                       ' Secretary-General Ban Ki-moon says his response to Russia\'s stepped up military support for'
                       ' Syria is that \'there is no military solution\' to the nearly five-year conflict and more weapons'
                       ' will only worsen the violence and misery for millions of people.',
                   ]
                   __lowerCamelCase   :	Dict								=		[
                       'Şeful ONU declară că nu există o soluţie militară în Siria',
                       'Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al'
                       ' Rusiei pentru Siria este că "nu există o soluţie militară" la conflictul de aproape cinci ani şi'
                       ' că noi arme nu vor face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.',
                   ]
                   try:
                        __lowerCamelCase   :	Any								=		tokenizer.prepare_seqaseq_batch(
                            src_texts=_lowerCAmelCase			,     tgt_texts=_lowerCAmelCase			,     max_length=3			,     max_target_length=10			,     return_tensors='pt'			,     src_lang='eng_Latn'			,     tgt_lang='ron_Latn'			,     )
                   except NotImplementedError:
                        return
                   self.assertEqual(batch.input_ids.shape[1]			,     3	)
                   self.assertEqual(batch.labels.shape[1]			,     10	)
                   # max_target_length will default to max_length if not specified
                   __lowerCamelCase   :	Optional[Any]								=		tokenizer.prepare_seqaseq_batch(
                       _lowerCAmelCase			,     tgt_texts=_lowerCAmelCase			,     max_length=3			,     return_tensors='pt'	)
                   self.assertEqual(batch.input_ids.shape[1]			,     3	)
                   self.assertEqual(batch.labels.shape[1]			,     3	)

                   __lowerCamelCase   :	str								=		tokenizer.prepare_seqaseq_batch(
                       src_texts=_lowerCAmelCase			,     max_length=3			,     max_target_length=10			,     return_tensors='pt'	)
                   self.assertEqual(batch_encoder_only.input_ids.shape[1]			,     3	)
                   self.assertEqual(batch_encoder_only.attention_mask.shape[1]			,     3	)
                   self.assertNotIn('decoder_input_ids'			,     _lowerCAmelCase	)







    @unittest.skip('Unfortunately way too slow to build a BPE with SentencePiece.'	)
    def       lowercase_				(							self	)				->     Union[str, Any]:
         pass







    def       lowercase_				(							self	)				->     Dict:
         for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
              with self.subTest(f'{tokenizer.__class__.__name__} ({pretrained_name})'	):
                   __lowerCamelCase   :	Tuple								=		[AddedToken('<special>'			,     lstrip=_lowerCAmelCase	)]

                   __lowerCamelCase   :	Any								=		self.rust_tokenizer_class.from_pretrained(
                       _lowerCAmelCase			,     additional_special_tokens=_lowerCAmelCase			,     **_lowerCAmelCase	)
                   __lowerCamelCase   :	Optional[int]								=		tokenizer_r.encode('Hey this is a <special> token'	)

                   __lowerCamelCase   :	List[Any]								=		tokenizer_r.encode('<special>'			,     add_special_tokens=_lowerCAmelCase	)[0]

                   self.assertTrue(special_token_id in r_output	)

                   if self.test_slow_tokenizer:
                        __lowerCamelCase   :	Union[str, Any]								=		self.rust_tokenizer_class.from_pretrained(
                            _lowerCAmelCase			,     additional_special_tokens=_lowerCAmelCase			,     **_lowerCAmelCase			,     )
                        __lowerCamelCase   :	Optional[Any]								=		self.tokenizer_class.from_pretrained(
                            _lowerCAmelCase			,     additional_special_tokens=_lowerCAmelCase			,     **_lowerCAmelCase	)

                        __lowerCamelCase   :	Optional[int]								=		tokenizer_p.encode('Hey this is a <special> token'	)

                        __lowerCamelCase   :	int								=		tokenizer_cr.encode('Hey this is a <special> token'	)

                        self.assertEqual(_lowerCAmelCase			,     _lowerCAmelCase	)
                        self.assertEqual(_lowerCAmelCase			,     _lowerCAmelCase	)
                        self.assertTrue(special_token_id in p_output	)
                        self.assertTrue(special_token_id in cr_output	)







@require_torch
@require_sentencepiece
@require_tokenizers
class    UpperCAmelCase_  (unittest.TestCase  ):





    """simple docstring"""






    lowerCamelCase  :					Optional[Any]							     =			'facebook/nllb-200-distilled-600M'
    lowerCamelCase  :					List[Any]							     =			[
        ' UN Chief Says There Is No Military Solution in Syria',
        ' Secretary-General Ban Ki-moon says his response to Russia\'s stepped up military support for Syria is that \"there is no military solution\" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.',
    ]
    lowerCamelCase  :					Optional[Any]							     =			[
        'Şeful ONU declară că nu există o soluţie militară în Siria',
        'Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei'
        ' pentru Siria este că \"nu există o soluţie militară\" la conflictul de aproape cinci ani şi că noi arme nu vor'
        ' face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.',
    ]
    lowerCamelCase  :					Dict							     =			[
        2_5_6_0_4_7,
        1_6_2_9_7,
        1_3_4_4_0_8,
        8_1_6_5,
        2_4_8_0_6_6,
        1_4_7_3_4,
        9_5_0,
        1_1_3_5,
        1_0_5_7_2_1,
        3_5_7_3,
        8_3,
        2_7_3_5_2,
        1_0_8,
        4_9_4_8_6,
        2,
    ]







    @classmethod
    def       lowercase_				(							cls	)				->     int:
         __lowerCamelCase   :	List[Any]								=		NllbTokenizer.from_pretrained(
             cls.checkpoint_name			,     src_lang='eng_Latn'			,     tgt_lang='ron_Latn'	)
         __lowerCamelCase   :	Optional[Any]								=		1
         return cls







    def       lowercase_				(							self	)				->     List[str]:
         self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ace_Arab']			,     25_60_01	)
         self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ace_Latn']			,     25_60_02	)
         self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['fra_Latn']			,     25_60_57	)







    def       lowercase_				(							self	)				->     Dict:
         __lowerCamelCase   :	List[Any]								=		self.tokenizer.batch_encode_plus(self.src_text	).input_ids[0]
         self.assertListEqual(self.expected_src_tokens			,     _lowerCAmelCase	)







    def       lowercase_				(							self	)				->     List[str]:
         self.assertIn(_lowerCAmelCase			,     self.tokenizer.all_special_ids	)
         # fmt: off
         __lowerCamelCase   :	List[str]								=		[RO_CODE, 42_54, 9_80_68, 11_29_23, 3_90_72, 39_09, 7_13, 10_27_67, 26, 1_73_14, 3_56_42, 1_46_83, 3_31_18, 20_22, 6_69_87, 2, 25_60_47]
         # fmt: on

         __lowerCamelCase   :	List[Any]								=		self.tokenizer.decode(_lowerCAmelCase			,     skip_special_tokens=_lowerCAmelCase	)
         __lowerCamelCase   :	int								=		self.tokenizer.decode(generated_ids[1:]			,     skip_special_tokens=_lowerCAmelCase	)
         self.assertEqual(_lowerCAmelCase			,     _lowerCAmelCase	)
         self.assertNotIn(self.tokenizer.eos_token			,     _lowerCAmelCase	)







    def       lowercase_				(							self	)				->     List[Any]:
         __lowerCamelCase   :	Any								=		['this is gunna be a long sentence ' * 20]
         assert isinstance(src_text[0]			,     _lowerCAmelCase	)
         __lowerCamelCase   :	Optional[int]								=		10
         __lowerCamelCase   :	Optional[int]								=		self.tokenizer(_lowerCAmelCase			,     max_length=_lowerCAmelCase			,     truncation=_lowerCAmelCase	).input_ids[0]
         self.assertEqual(ids[-1]			,     2	)
         self.assertEqual(ids[0]			,     _lowerCAmelCase	)
         self.assertEqual(len(_lowerCAmelCase	)			,     _lowerCAmelCase	)







    def       lowercase_				(							self	)				->     Optional[int]:
         self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['<mask>', 'ar_AR']	)			,     [25_62_03, 3]	)







    def       lowercase_				(							self	)				->     Optional[int]:
         __lowerCamelCase   :	Any								=		tempfile.mkdtemp()
         __lowerCamelCase   :	Optional[Any]								=		self.tokenizer.fairseq_tokens_to_ids
         self.tokenizer.save_pretrained(_lowerCAmelCase	)
         __lowerCamelCase   :	int								=		NllbTokenizer.from_pretrained(_lowerCAmelCase	)
         self.assertDictEqual(new_tok.fairseq_tokens_to_ids			,     _lowerCAmelCase	)







    @require_torch
    def       lowercase_				(							self	)				->     Optional[Any]:
         __lowerCamelCase   :	Tuple								=		self.tokenizer(
             self.src_text			,     text_target=self.tgt_text			,     padding=_lowerCAmelCase			,     truncation=_lowerCAmelCase			,     max_length=len(self.expected_src_tokens	)			,     return_tensors='pt'			,     )
         __lowerCamelCase   :	Union[str, Any]								=		shift_tokens_right(
             batch['labels']			,     self.tokenizer.pad_token_id			,     self.tokenizer.lang_code_to_id['ron_Latn']	)

         self.assertIsInstance(_lowerCAmelCase			,     _lowerCAmelCase	)

         self.assertEqual((2, 15)			,     batch.input_ids.shape	)
         self.assertEqual((2, 15)			,     batch.attention_mask.shape	)
         __lowerCamelCase   :	Union[str, Any]								=		batch.input_ids.tolist()[0]
         self.assertListEqual(self.expected_src_tokens			,     _lowerCAmelCase	)
         self.assertEqual(_lowerCAmelCase			,     batch.decoder_input_ids[0, 0]	)  # EOS
         # Test that special tokens are reset
         self.assertEqual(self.tokenizer.prefix_tokens			,     [EN_CODE]	)
         self.assertEqual(self.tokenizer.suffix_tokens			,     [self.tokenizer.eos_token_id]	)







    def       lowercase_				(							self	)				->     str:
         __lowerCamelCase   :	List[Any]								=		self.tokenizer(self.src_text			,     padding=_lowerCAmelCase			,     truncation=_lowerCAmelCase			,     max_length=3			,     return_tensors='pt'	)
         __lowerCamelCase   :	int								=		self.tokenizer(
             text_target=self.tgt_text			,     padding=_lowerCAmelCase			,     truncation=_lowerCAmelCase			,     max_length=10			,     return_tensors='pt'	)
         __lowerCamelCase   :	int								=		targets['input_ids']
         __lowerCamelCase   :	List[Any]								=		shift_tokens_right(
             _lowerCAmelCase			,     self.tokenizer.pad_token_id			,     decoder_start_token_id=self.tokenizer.lang_code_to_id[self.tokenizer.tgt_lang]			,     )

         self.assertEqual(batch.input_ids.shape[1]			,     3	)
         self.assertEqual(batch.decoder_input_ids.shape[1]			,     10	)







    @require_torch
    def       lowercase_				(							self	)				->     Tuple:
         __lowerCamelCase   :	Tuple								=		self.tokenizer._build_translation_inputs(
             'A test'			,     return_tensors='pt'			,     src_lang='eng_Latn'			,     tgt_lang='fra_Latn'	)

         self.assertEqual(
             nested_simplify(_lowerCAmelCase	)			,     {
                 # A, test, EOS, en_XX
                 'input_ids': [[25_60_47, 70, 73_56, 2]],
                 'attention_mask': [[1, 1, 1, 1]],
                 # ar_AR
                 'forced_bos_token_id': 25_60_57,
             }			,     )







    @require_torch
    def       lowercase_				(							self	)				->     Union[str, Any]:
         __lowerCamelCase   :	Dict								=		True
         __lowerCamelCase   :	Tuple								=		self.tokenizer(
             'UN Chief says there is no military solution in Syria'			,     src_lang='eng_Latn'			,     tgt_lang='fra_Latn'	)
         self.assertEqual(
             inputs.input_ids			,     [1_62_97, 13_44_08, 2_56_53, 63_70, 2_48, 2_54, 10_39_29, 9_49_95, 1_08, 4_94_86, 2, 25_60_47]	)

         __lowerCamelCase   :	str								=		False
         __lowerCamelCase   :	str								=		self.tokenizer(
             'UN Chief says there is no military solution in Syria'			,     src_lang='eng_Latn'			,     tgt_lang='fra_Latn'	)
         self.assertEqual(
             inputs.input_ids			,     [25_60_47, 1_62_97, 13_44_08, 2_56_53, 63_70, 2_48, 2_54, 10_39_29, 9_49_95, 1_08, 4_94_86, 2]	)












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

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	List[Any]							)					-> Union[str, Any]:
      SCREAMING_SNAKE_CASE_							=			['head.weight', 'head.bias']
      for k in ignore_keys:
            state_dict.pop(__UpperCAmelCase			,			__UpperCAmelCase							)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Optional[int]							)					-> str:
      # projection head is used in the self-supervised pre-training in MSN,
      # for downstream task it's not needed.
      SCREAMING_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(__UpperCAmelCase			,			__UpperCAmelCase							)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Dict			,			__UpperCAmelCase						:	Any			,			__UpperCAmelCase						:	Any							)					-> Dict:
      SCREAMING_SNAKE_CASE_							=			dct.pop(__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_							=			val

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	Optional[int]			,			__UpperCAmelCase						:	Tuple							)					-> List[str]:
      SCREAMING_SNAKE_CASE_							=			ViTMSNConfig()
      SCREAMING_SNAKE_CASE_							=			10_00

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

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

      SCREAMING_SNAKE_CASE_							=			ViTMSNModel(__UpperCAmelCase							)

      SCREAMING_SNAKE_CASE_							=			torch.hub.load_state_dict_from_url(__UpperCAmelCase			,			map_location='cpu'							)['target_encoder']

      SCREAMING_SNAKE_CASE_							=			ViTImageProcessor(size=config.image_size							)

      remove_projection_head(__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_							=			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()

      SCREAMING_SNAKE_CASE_							=			'http://images.cocodataset.org/val2017/000000039769.jpg'

      SCREAMING_SNAKE_CASE_							=			Image.open(requests.get(__UpperCAmelCase			,			stream=__UpperCAmelCase							).raw							)
      SCREAMING_SNAKE_CASE_							=			ViTImageProcessor(
          size=config.image_size			,			image_mean=__UpperCAmelCase			,			image_std=__UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_							=			image_processor(images=__UpperCAmelCase			,			return_tensors='pt'							)

      # forward pass
      torch.manual_seed(2							)
      SCREAMING_SNAKE_CASE_							=			model(**__UpperCAmelCase							)
      SCREAMING_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:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]]							)
      elif "b16" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[1_4.2_8_8_9, -1_8.9_0_4_5, 1_1.7_2_8_1]]							)
      elif "l16" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[4_1.5_0_2_8, -2_2.8_6_8_1, 4_5.6_4_7_5]]							)
      elif "b4" in checkpoint_url:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]]							)
      else:
            SCREAMING_SNAKE_CASE_							=			torch.tensor([[-0.1_7_9_2, -0.6_4_6_5, 2.4_2_6_3]]							)

      # 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__ :   Union[str, Any]					=							parser.parse_args()
   convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)