

'''simple docstring'''






import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
     DistilBertForMaskedLM,
     DistilBertForMultipleChoice,
     DistilBertForQuestionAnswering,
     DistilBertForSequenceClassification,
     DistilBertForTokenClassification,
     DistilBertModel,
 )






class 			lowercase_	(lowerCamelCase__      ):




  """simple docstring"""




  def __init__(      self   :     Optional[Any]					,lowercase__   :     Union[str, Any]					,lowercase__   :     int=1_3					,lowercase__   :     Optional[int]=7					,lowercase__   :     List[Any]=True					,lowercase__   :     Tuple=True					,lowercase__   :     List[Any]=False					,lowercase__   :     Dict=True					,lowercase__   :     List[str]=9_9					,lowercase__   :     List[str]=3_2					,lowercase__   :     Union[str, Any]=5					,lowercase__   :     int=4					,lowercase__   :     Tuple=3_7					,lowercase__   :     List[str]="gelu"					,lowercase__   :     Dict=0.1					,lowercase__   :     List[Any]=0.1					,lowercase__   :     Any=5_1_2					,lowercase__   :     Tuple=1_6					,lowercase__   :     Optional[Any]=2					,lowercase__   :     str=0.0_2					,lowercase__   :     List[Any]=3					,lowercase__   :     Any=4					,lowercase__   :     List[Any]=None					,):
   __lowercase							  = parent
   __lowercase							  = batch_size
   __lowercase							  = seq_length
   __lowercase							  = is_training
   __lowercase							  = use_input_mask
   __lowercase							  = use_token_type_ids
   __lowercase							  = use_labels
   __lowercase							  = vocab_size
   __lowercase							  = hidden_size
   __lowercase							  = num_hidden_layers
   __lowercase							  = num_attention_heads
   __lowercase							  = intermediate_size
   __lowercase							  = hidden_act
   __lowercase							  = hidden_dropout_prob
   __lowercase							  = attention_probs_dropout_prob
   __lowercase							  = max_position_embeddings
   __lowercase							  = type_vocab_size
   __lowercase							  = type_sequence_label_size
   __lowercase							  = initializer_range
   __lowercase							  = num_labels
   __lowercase							  = num_choices
   __lowercase							  = scope






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   __lowercase							  = ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size			)

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

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

   __lowercase							  = self.get_config()

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






  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   return DistilBertConfig(
       vocab_size=self.vocab_size					,dim=self.hidden_size					,n_layers=self.num_hidden_layers					,n_heads=self.num_attention_heads					,hidden_dim=self.intermediate_size					,hidden_act=self.hidden_act					,dropout=self.hidden_dropout_prob					,attention_dropout=self.attention_probs_dropout_prob					,max_position_embeddings=self.max_position_embeddings					,initializer_range=self.initializer_range					,)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict					,lowercase__   :     Optional[int]					,lowercase__   :     Tuple					,lowercase__   :     Tuple					,lowercase__   :     Union[str, Any]					,lowercase__   :     Any					,lowercase__   :     int			):
   __lowercase							  = DistilBertModel(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,lowercase__			)
   __lowercase							  = model(lowercase__			)
   self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]					,lowercase__   :     Optional[int]					,lowercase__   :     Optional[Any]					,lowercase__   :     Tuple					,lowercase__   :     Optional[int]					,lowercase__   :     List[Any]					,lowercase__   :     Tuple			):
   __lowercase							  = DistilBertForMaskedLM(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.vocab_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]					,lowercase__   :     Tuple					,lowercase__   :     int					,lowercase__   :     List[Any]					,lowercase__   :     int					,lowercase__   :     Dict					,lowercase__   :     Union[str, Any]			):
   __lowercase							  = DistilBertForQuestionAnswering(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(
       lowercase__					,attention_mask=lowercase__					,start_positions=lowercase__					,end_positions=lowercase__			)
   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 					SCREAMING_SNAKE_CASE      (      self   :     Tuple					,lowercase__   :     Dict					,lowercase__   :     Tuple					,lowercase__   :     str					,lowercase__   :     Dict					,lowercase__   :     Optional[int]					,lowercase__   :     Tuple			):
   __lowercase							  = self.num_labels
   __lowercase							  = DistilBertForSequenceClassification(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_labels)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     List[Any]					,lowercase__   :     Tuple					,lowercase__   :     Union[str, Any]					,lowercase__   :     Optional[int]			):
   __lowercase							  = self.num_labels
   __lowercase							  = DistilBertForTokenClassification(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()

   __lowercase							  = model(lowercase__					,attention_mask=lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.seq_length, self.num_labels)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]					,lowercase__   :     List[str]					,lowercase__   :     Dict					,lowercase__   :     Tuple					,lowercase__   :     List[Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any			):
   __lowercase							  = self.num_choices
   __lowercase							  = DistilBertForMultipleChoice(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = input_ids.unsqueeze(1			).expand(-1					,self.num_choices					,-1			).contiguous()
   __lowercase							  = input_mask.unsqueeze(1			).expand(-1					,self.num_choices					,-1			).contiguous()
   __lowercase							  = model(
       lowercase__					,attention_mask=lowercase__					,labels=lowercase__					,)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_choices)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase							  = self.prepare_config_and_inputs()
   ((__lowercase)  ,	(__lowercase)  ,	(__lowercase)  ,	(__lowercase)  ,	(__lowercase)  ,	(__lowercase))							  = config_and_inputs
   __lowercase							  = {'''input_ids''': input_ids, '''attention_mask''': input_mask}
   return config, inputs_dict






@require_torch
class 			lowercase_	(lowerCamelCase__				,       lowerCamelCase__				,       unittest.TestCase      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  List[Any]							=			(
      (
          DistilBertModel,
          DistilBertForMaskedLM,
          DistilBertForMultipleChoice,
          DistilBertForQuestionAnswering,
          DistilBertForSequenceClassification,
          DistilBertForTokenClassification,
      )
      if is_torch_available()
      else None
  )
  SCREAMING_SNAKE_CASE    :  Dict							=			(
      {
          'feature-extraction': DistilBertModel,
          'fill-mask': DistilBertForMaskedLM,
          'question-answering': DistilBertForQuestionAnswering,
          'text-classification': DistilBertForSequenceClassification,
          'token-classification': DistilBertForTokenClassification,
          'zero-shot': DistilBertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  SCREAMING_SNAKE_CASE    :  Union[str, Any]							=			True
  SCREAMING_SNAKE_CASE    :  Dict							=			True
  SCREAMING_SNAKE_CASE    :  Optional[Any]							=			True
  SCREAMING_SNAKE_CASE    :  int							=			True






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = DistilBertModelTester(self			)
   __lowercase							  = ConfigTester(self					,config_class=lowercase__					,dim=3_7			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   self.config_tester.run_common_tests()






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_model(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_masked_lm(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_question_answering(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_token_classification(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowercase__			)






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict			):
   for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowercase							  = DistilBertModel.from_pretrained(lowercase__			)
    self.assertIsNotNone(lowercase__			)






  @slow
  @require_torch_gpu
  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()
   for model_class in self.all_model_classes:
    # BertForMultipleChoice behaves incorrectly in JIT environments.
    if model_class == DistilBertForMultipleChoice:
     return

    __lowercase							  = True
    __lowercase							  = model_class(config=lowercase__			)

    __lowercase							  = self._prepare_for_class(lowercase__					,lowercase__			)
    __lowercase							  = torch.jit.trace(
        lowercase__					,(inputs_dict['''input_ids'''].to('''cpu'''			), inputs_dict['''attention_mask'''].to('''cpu'''			))			)

    with tempfile.TemporaryDirectory() as tmp:
     torch.jit.save(lowercase__					,os.path.join(lowercase__					,'''traced_model.pt'''			)			)
     __lowercase							  = torch.jit.load(os.path.join(lowercase__					,'''traced_model.pt'''			)					,map_location=lowercase__			)
     loaded(inputs_dict['''input_ids'''].to(lowercase__			)					,inputs_dict['''attention_mask'''].to(lowercase__			)			)






@require_torch
class 			lowercase_	(unittest.TestCase      ):




  """simple docstring"""




  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase							  = DistilBertModel.from_pretrained('''distilbert-base-uncased'''			)
   __lowercase							  = torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]			)
   __lowercase							  = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
   with torch.no_grad():
    __lowercase							  = model(lowercase__					,attention_mask=lowercase__			)[0]
   __lowercase							  = torch.Size((1, 1_1, 7_6_8)			)
   self.assertEqual(output.shape					,lowercase__			)
   __lowercase							  = torch.tensor(
       [[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]]			)

   self.assertTrue(torch.allclose(output[:, 1:4, 1:4]					,lowercase__					,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)