




from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import (
       TF_MODEL_FOR_PRETRAINING_MAPPING,
       TFMobileBertForMaskedLM,
       TFMobileBertForMultipleChoice,
       TFMobileBertForNextSentencePrediction,
       TFMobileBertForPreTraining,
       TFMobileBertForQuestionAnswering,
       TFMobileBertForSequenceClassification,
       TFMobileBertForTokenClassification,
       TFMobileBertModel,
   )





@require_tf
class        lowerCamelCase_							(      _SCREAMING_SNAKE_CASE       ,   _SCREAMING_SNAKE_CASE       ,   unittest.TestCase					):


 '''simple docstring'''







 lowercase_				    =    (
     (
         TFMobileBertModel,
         TFMobileBertForMaskedLM,
         TFMobileBertForNextSentencePrediction,
         TFMobileBertForPreTraining,
         TFMobileBertForQuestionAnswering,
         TFMobileBertForSequenceClassification,
         TFMobileBertForTokenClassification,
         TFMobileBertForMultipleChoice,
     )
     if is_tf_available()
     else ()
 )
 lowercase_				    =    (
     {
         "feature-extraction": TFMobileBertModel,
         "fill-mask": TFMobileBertForMaskedLM,
         "question-answering": TFMobileBertForQuestionAnswering,
         "text-classification": TFMobileBertForSequenceClassification,
         "token-classification": TFMobileBertForTokenClassification,
         "zero-shot": TFMobileBertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 lowercase_				    =    False
 lowercase_				    =    False



 def       lowerCAmelCase_  (							self  :		Union[str, Any]				,				_lowerCAmelCase  :		int				,				_lowerCAmelCase  :		str				,				_lowerCAmelCase  :		Optional[int]=False						):
       SCREAMING_SNAKE_CASE_							=			super()._prepare_for_class(_lowerCAmelCase				,				_lowerCAmelCase				,				return_labels=_lowerCAmelCase						)

       if return_labels:
             if model_class in get_values(_lowerCAmelCase						):
                   SCREAMING_SNAKE_CASE_							=			tf.zeros(self.model_tester.batch_size				,				dtype=tf.intaa						)

       return inputs_dict





 class        lowerCamelCase_							(      _SCREAMING_SNAKE_CASE					):


  '''simple docstring'''







  def __init__(							self  :		int				,				_lowerCAmelCase  :		Dict				,				_lowerCAmelCase  :		str=13				,				_lowerCAmelCase  :		str=7				,				_lowerCAmelCase  :		Optional[Any]=True				,				_lowerCAmelCase  :		Optional[int]=True				,				_lowerCAmelCase  :		Optional[Any]=True				,				_lowerCAmelCase  :		Union[str, Any]=True				,				_lowerCAmelCase  :		Union[str, Any]=99				,				_lowerCAmelCase  :		str=32				,				_lowerCAmelCase  :		Dict=32				,				_lowerCAmelCase  :		str=2				,				_lowerCAmelCase  :		Optional[Any]=4				,				_lowerCAmelCase  :		List[Any]=37				,				_lowerCAmelCase  :		Any="gelu"				,				_lowerCAmelCase  :		List[str]=0.1				,				_lowerCAmelCase  :		int=0.1				,				_lowerCAmelCase  :		List[str]=512				,				_lowerCAmelCase  :		List[Any]=16				,				_lowerCAmelCase  :		Tuple=2				,				_lowerCAmelCase  :		Optional[Any]=0.02				,				_lowerCAmelCase  :		Any=3				,				_lowerCAmelCase  :		Optional[int]=4				,				_lowerCAmelCase  :		str=None				,				):
        SCREAMING_SNAKE_CASE_							=			parent
        SCREAMING_SNAKE_CASE_							=			batch_size
        SCREAMING_SNAKE_CASE_							=			seq_length
        SCREAMING_SNAKE_CASE_							=			is_training
        SCREAMING_SNAKE_CASE_							=			use_input_mask
        SCREAMING_SNAKE_CASE_							=			use_token_type_ids
        SCREAMING_SNAKE_CASE_							=			use_labels
        SCREAMING_SNAKE_CASE_							=			vocab_size
        SCREAMING_SNAKE_CASE_							=			hidden_size
        SCREAMING_SNAKE_CASE_							=			num_hidden_layers
        SCREAMING_SNAKE_CASE_							=			num_attention_heads
        SCREAMING_SNAKE_CASE_							=			intermediate_size
        SCREAMING_SNAKE_CASE_							=			hidden_act
        SCREAMING_SNAKE_CASE_							=			hidden_dropout_prob
        SCREAMING_SNAKE_CASE_							=			attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE_							=			max_position_embeddings
        SCREAMING_SNAKE_CASE_							=			type_vocab_size
        SCREAMING_SNAKE_CASE_							=			type_sequence_label_size
        SCREAMING_SNAKE_CASE_							=			initializer_range
        SCREAMING_SNAKE_CASE_							=			num_labels
        SCREAMING_SNAKE_CASE_							=			num_choices
        SCREAMING_SNAKE_CASE_							=			scope
        SCREAMING_SNAKE_CASE_							=			embedding_size



  def       lowerCAmelCase_  (							self  :		int						):
        SCREAMING_SNAKE_CASE_							=			ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size						)

        SCREAMING_SNAKE_CASE_							=			None
        if self.use_input_mask:
              SCREAMING_SNAKE_CASE_							=			random_attention_mask([self.batch_size, self.seq_length]						)

        SCREAMING_SNAKE_CASE_							=			None
        if self.use_token_type_ids:
              SCREAMING_SNAKE_CASE_							=			ids_tensor([self.batch_size, self.seq_length]				,				self.type_vocab_size						)

        SCREAMING_SNAKE_CASE_							=			None
        SCREAMING_SNAKE_CASE_							=			None
        SCREAMING_SNAKE_CASE_							=			None
        if self.use_labels:
              SCREAMING_SNAKE_CASE_							=			ids_tensor([self.batch_size]				,				self.type_sequence_label_size						)
              SCREAMING_SNAKE_CASE_							=			ids_tensor([self.batch_size, self.seq_length]				,				self.num_labels						)
              SCREAMING_SNAKE_CASE_							=			ids_tensor([self.batch_size]				,				self.num_choices						)

        SCREAMING_SNAKE_CASE_							=			MobileBertConfig(
            vocab_size=self.vocab_size				,				hidden_size=self.hidden_size				,				num_hidden_layers=self.num_hidden_layers				,				num_attention_heads=self.num_attention_heads				,				intermediate_size=self.intermediate_size				,				hidden_act=self.hidden_act				,				hidden_dropout_prob=self.hidden_dropout_prob				,				attention_probs_dropout_prob=self.attention_probs_dropout_prob				,				max_position_embeddings=self.max_position_embeddings				,				type_vocab_size=self.type_vocab_size				,				initializer_range=self.initializer_range				,				embedding_size=self.embedding_size				,				)

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



  def       lowerCAmelCase_  (							self  :		int				,				_lowerCAmelCase  :		List[Any]				,				_lowerCAmelCase  :		str				,				_lowerCAmelCase  :		Any				,				_lowerCAmelCase  :		str				,				_lowerCAmelCase  :		List[Any]				,				_lowerCAmelCase  :		List[str]				,				_lowerCAmelCase  :		List[str]						):
        SCREAMING_SNAKE_CASE_							=			TFMobileBertModel(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)

        SCREAMING_SNAKE_CASE_							=			[input_ids, input_mask]
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)

        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)

        self.parent.assertEqual(
            result.last_hidden_state.shape				,				(self.batch_size, self.seq_length, self.hidden_size)						)
        self.parent.assertEqual(result.pooler_output.shape				,				(self.batch_size, self.hidden_size)						)



  def       lowerCAmelCase_  (							self  :		Optional[Any]				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		Dict				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		int				,				_lowerCAmelCase  :		List[Any]						):
        SCREAMING_SNAKE_CASE_							=			TFMobileBertForMaskedLM(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.vocab_size)						)



  def       lowerCAmelCase_  (							self  :		str				,				_lowerCAmelCase  :		Optional[Any]				,				_lowerCAmelCase  :		Optional[Any]				,				_lowerCAmelCase  :		Optional[Any]				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		Any				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		Any						):
        SCREAMING_SNAKE_CASE_							=			TFMobileBertForNextSentencePrediction(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, 2)						)



  def       lowerCAmelCase_  (							self  :		Optional[int]				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		List[str]				,				_lowerCAmelCase  :		List[Any]				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		Any				,				_lowerCAmelCase  :		List[Any]				,				_lowerCAmelCase  :		Tuple						):
        SCREAMING_SNAKE_CASE_							=			TFMobileBertForPreTraining(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)
        self.parent.assertEqual(
            result.prediction_logits.shape				,				(self.batch_size, self.seq_length, self.vocab_size)						)
        self.parent.assertEqual(result.seq_relationship_logits.shape				,				(self.batch_size, 2)						)



  def       lowerCAmelCase_  (							self  :		Any				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		List[Any]				,				_lowerCAmelCase  :		List[str]				,				_lowerCAmelCase  :		Dict				,				_lowerCAmelCase  :		Any				,				_lowerCAmelCase  :		Optional[Any]				,				_lowerCAmelCase  :		List[str]						):
        SCREAMING_SNAKE_CASE_							=			self.num_labels
        SCREAMING_SNAKE_CASE_							=			TFMobileBertForSequenceClassification(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.num_labels)						)



  def       lowerCAmelCase_  (							self  :		str				,				_lowerCAmelCase  :		str				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		List[Any]				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		List[str]				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		Optional[Any]						):
        SCREAMING_SNAKE_CASE_							=			self.num_choices
        SCREAMING_SNAKE_CASE_							=			TFMobileBertForMultipleChoice(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			tf.tile(tf.expand_dims(_lowerCAmelCase				,				1						)				,				(1, self.num_choices, 1)						)
        SCREAMING_SNAKE_CASE_							=			tf.tile(tf.expand_dims(_lowerCAmelCase				,				1						)				,				(1, self.num_choices, 1)						)
        SCREAMING_SNAKE_CASE_							=			tf.tile(tf.expand_dims(_lowerCAmelCase				,				1						)				,				(1, self.num_choices, 1)						)
        SCREAMING_SNAKE_CASE_							=			{
            'input_ids': multiple_choice_inputs_ids,
            'attention_mask': multiple_choice_input_mask,
            'token_type_ids': multiple_choice_token_type_ids,
        }
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.num_choices)						)



  def       lowerCAmelCase_  (							self  :		int				,				_lowerCAmelCase  :		List[Any]				,				_lowerCAmelCase  :		Optional[int]				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		int				,				_lowerCAmelCase  :		int						):
        SCREAMING_SNAKE_CASE_							=			self.num_labels
        SCREAMING_SNAKE_CASE_							=			TFMobileBertForTokenClassification(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.num_labels)						)



  def       lowerCAmelCase_  (							self  :		Dict				,				_lowerCAmelCase  :		Dict				,				_lowerCAmelCase  :		Tuple				,				_lowerCAmelCase  :		List[str]				,				_lowerCAmelCase  :		int				,				_lowerCAmelCase  :		Any				,				_lowerCAmelCase  :		Optional[Any]				,				_lowerCAmelCase  :		Optional[int]						):
        SCREAMING_SNAKE_CASE_							=			TFMobileBertForQuestionAnswering(config=_lowerCAmelCase						)
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)
        self.parent.assertEqual(result.start_logits.shape				,				(self.batch_size, self.seq_length)						)
        self.parent.assertEqual(result.end_logits.shape				,				(self.batch_size, self.seq_length)						)




  def       lowerCAmelCase_  (							self  :		str						):
        SCREAMING_SNAKE_CASE_							=			self.prepare_config_and_inputs()
        (
            (
            SCREAMING_SNAKE_CASE_
        )						, (
            SCREAMING_SNAKE_CASE_
        )						, (
            SCREAMING_SNAKE_CASE_
        )						, (
            SCREAMING_SNAKE_CASE_
        )						, (
            SCREAMING_SNAKE_CASE_
        )						, (
            SCREAMING_SNAKE_CASE_
        )						, (
            SCREAMING_SNAKE_CASE_
        )						, 
        )							=			config_and_inputs
        SCREAMING_SNAKE_CASE_							=			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
        return config, inputs_dict



 def       lowerCAmelCase_  (							self  :		Dict						):
       SCREAMING_SNAKE_CASE_							=			TFMobileBertModelTest.TFMobileBertModelTester(self						)
       SCREAMING_SNAKE_CASE_							=			ConfigTester(self				,				config_class=_lowerCAmelCase				,				hidden_size=37						)



 def       lowerCAmelCase_  (							self  :		Optional[Any]						):
       self.config_tester.run_common_tests()



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_model(*_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		List[Any]						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_masked_lm(*_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_multiple_choice(*_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		Dict						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_pretraining(*_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_question_answering(*_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		Union[str, Any]						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_sequence_classification(*_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mobilebert_for_token_classification(*_lowerCAmelCase						)




 @slow
 def       lowerCAmelCase_  (							self  :		Any						):
       # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       for model_name in ["google/mobilebert-uncased"]:
             SCREAMING_SNAKE_CASE_							=			TFMobileBertModel.from_pretrained(_lowerCAmelCase						)
             self.assertIsNotNone(_lowerCAmelCase						)





@require_tf
class        lowerCamelCase_							(      unittest.TestCase					):


 '''simple docstring'''







 @slow
 def       lowerCAmelCase_  (							self  :		Tuple						):
       SCREAMING_SNAKE_CASE_							=			TFMobileBertForPreTraining.from_pretrained('google/mobilebert-uncased'						)
       SCREAMING_SNAKE_CASE_							=			tf.constant([[0, 1, 2, 3, 4, 5]]						)
       SCREAMING_SNAKE_CASE_							=			model(_lowerCAmelCase						)[0]

       SCREAMING_SNAKE_CASE_							=			[1, 6, 30_522]
       self.assertEqual(output.shape				,				_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_							=			tf.constant(
           [
               [
                   [-4.591_9547, -9.24_8295, -9.64_5256],
                   [-6.730_6175, -6.44_0284, -6.605_2837],
                   [-7.274_3506, -6.784_7915, -6.02_4673],
               ]
           ]						)
       tf.debugging.assert_near(output[:, :3, :3]				,				_lowerCAmelCase				,				atol=1E-4						)





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)