


from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
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 (
     TFRoFormerForCausalLM,
     TFRoFormerForMaskedLM,
     TFRoFormerForMultipleChoice,
     TFRoFormerForQuestionAnswering,
     TFRoFormerForSequenceClassification,
     TFRoFormerForTokenClassification,
     TFRoFormerModel,
 )
 from transformers.models.roformer.modeling_tf_roformer import (
     TFRoFormerSelfAttention,
     TFRoFormerSinusoidalPositionalEmbedding,
 )

class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case=1_3      , __snake_case=7      , __snake_case=True      , __snake_case=True      , __snake_case=True      , __snake_case=True      , __snake_case=9_9      , __snake_case=3_2      , __snake_case=2      , __snake_case=4      , __snake_case=3_7      , __snake_case="gelu"      , __snake_case=0.1      , __snake_case=0.1      , __snake_case=5_1_2      , __snake_case=1_6      , __snake_case=2      , __snake_case=0.02      , __snake_case=3      , __snake_case=4      , __snake_case=None      , ):
         snake_case					=    parent
         snake_case					=    1_3
         snake_case					=    7
         snake_case					=    True
         snake_case					=    True
         snake_case					=    True
         snake_case					=    True
         snake_case					=    9_9
         snake_case					=    3_2
         snake_case					=    2
         snake_case					=    4
         snake_case					=    3_7
         snake_case					=    '''gelu'''
         snake_case					=    0.1
         snake_case					=    0.1
         snake_case					=    5_1_2
         snake_case					=    1_6
         snake_case					=    2
         snake_case					=    0.02
         snake_case					=    3
         snake_case					=    4
         snake_case					=    None
    def 		a_      (	self   ):
         snake_case					=    ids_tensor([self.batch_size, self.seq_length]      , self.vocab_size   )

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

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

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

         snake_case					=    RoFormerConfig(
             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      , return_dict=__snake_case      , )

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    TFRoFormerModel(config=__snake_case   )
         snake_case					=    {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

         snake_case					=    [input_ids, input_mask]
         snake_case					=    model(__snake_case   )

         snake_case					=    model(__snake_case   )

         self.parent.assertEqual(result.last_hidden_state.shape      , (self.batch_size, self.seq_length, self.hidden_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    True
         snake_case					=    TFRoFormerForCausalLM(config=__snake_case   )
         snake_case					=    {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }
         snake_case					=    model(__snake_case   )['''logits''']
         self.parent.assertListEqual(
             list(prediction_scores.numpy().shape   )      , [self.batch_size, self.seq_length, self.vocab_size]   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    TFRoFormerForMaskedLM(config=__snake_case   )
         snake_case					=    {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }
         snake_case					=    model(__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.vocab_size)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    TFRoFormerForSequenceClassification(config=__snake_case   )
         snake_case					=    {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }

         snake_case					=    model(__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.num_labels)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_choices
         snake_case					=    TFRoFormerForMultipleChoice(config=__snake_case   )
         snake_case					=    tf.tile(tf.expand_dims(__snake_case      , 1   )      , (1, self.num_choices, 1)   )
         snake_case					=    tf.tile(tf.expand_dims(__snake_case      , 1   )      , (1, self.num_choices, 1)   )
         snake_case					=    tf.tile(tf.expand_dims(__snake_case      , 1   )      , (1, self.num_choices, 1)   )
         snake_case					=    {
             '''input_ids''': multiple_choice_inputs_ids,
             '''attention_mask''': multiple_choice_input_mask,
             '''token_type_ids''': multiple_choice_token_type_ids,
         }
         snake_case					=    model(__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.num_choices)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    self.num_labels
         snake_case					=    TFRoFormerForTokenClassification(config=__snake_case   )
         snake_case					=    {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }
         snake_case					=    model(__snake_case   )
         self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.seq_length, self.num_labels)   )
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    TFRoFormerForQuestionAnswering(config=__snake_case   )
         snake_case					=    {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }

         snake_case					=    model(__snake_case   )
         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 		a_      (	self   ):
         snake_case					=    self.prepare_config_and_inputs()
         (
             (
             snake_case
         )			,  (
             snake_case
         )			,  (
             snake_case
         )			,  (
             snake_case
         )			,  (
             snake_case
         )			,  (
             snake_case
         )			,  (
             snake_case
         )			,  
         )					=    config_and_inputs
         snake_case					=    {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
         return config, inputs_dict






@require_tf
class   A__ (   snake_case__	, snake_case__	, unittest.TestCase     ):







    """simple docstring"""

    __magic_name__        =		(
        (
            TFRoFormerModel,
            TFRoFormerForCausalLM,
            TFRoFormerForMaskedLM,
            TFRoFormerForQuestionAnswering,
            TFRoFormerForSequenceClassification,
            TFRoFormerForTokenClassification,
            TFRoFormerForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    __magic_name__        =		(
        {
            'feature-extraction': TFRoFormerModel,
            'fill-mask': TFRoFormerForMaskedLM,
            'question-answering': TFRoFormerForQuestionAnswering,
            'text-classification': TFRoFormerForSequenceClassification,
            'text-generation': TFRoFormerForCausalLM,
            'token-classification': TFRoFormerForTokenClassification,
            'zero-shot': TFRoFormerForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )

    __magic_name__        =		False
    __magic_name__        =		False
    def 		a_      (	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         if pipeline_test_casse_name == "TextGenerationPipelineTests":
              return True

         return False
    def 		a_      (	self   ):
         snake_case					=    TFRoFormerModelTester(self   )
         snake_case					=    ConfigTester(self      , config_class=__snake_case      , hidden_size=3_7   )
    def 		a_      (	self   ):
         self.config_tester.run_common_tests()
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__snake_case   )
    def 		a_      (	self   ):
         snake_case					=    self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__snake_case   )




    @slow
    def 		a_      (	self   ):
         snake_case					=    TFRoFormerModel.from_pretrained('''junnyu/roformer_chinese_base'''   )
         self.assertIsNotNone(__snake_case   )






@require_tf
class   A__ (   unittest.TestCase     ):







    """simple docstring"""

    @slow
    def 		a_      (	self   ):
         snake_case					=    TFRoFormerForMaskedLM.from_pretrained('''junnyu/roformer_chinese_base'''   )
         snake_case					=    tf.constant([[0, 1, 2, 3, 4, 5]]   )
         snake_case					=    model(__snake_case   )[0]

         # TODO Replace vocab size
         snake_case					=    5_0_0_0_0

         snake_case					=    [1, 6, vocab_size]
         self.assertEqual(output.shape      , __snake_case   )

         print(output[:, :3, :3]   )

         # TODO Replace values below with what was printed above.
         snake_case					=    tf.constant(
             [
                 [
                     [-0.1205_3341, -1.026_4901, 0.2922_1946],
                     [-1.513_3783, 0.19_7433, 0.1519_0607],
                     [-5.013_5403, -3.90_0256, -0.8403_8764],
                 ]
             ]   )
         tf.debugging.assert_near(output[:, :3, :3]      , __snake_case      , atol=1E-4   )






@require_tf
class   A__ (   unittest.TestCase     ):







    """simple docstring"""

    __magic_name__        =		1e-4
    def 		a_      (	self   ):
         snake_case					=    tf.constant([[4, 1_0]]   )
         snake_case					=    TFRoFormerSinusoidalPositionalEmbedding(num_positions=6      , embedding_dim=6   )

         snake_case					=    emba(input_ids.shape   )
         snake_case					=    tf.constant(
             [[0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 1.0000], [0.8415, 0.0464, 0.0022, 0.5403, 0.9989, 1.0000]]   )

         tf.debugging.assert_near(__snake_case      , __snake_case      , atol=self.tolerance   )




    def 		a_      (	self   ):
         snake_case					=    tf.constant(
             [
                 [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                 [0.8415, 0.8219, 0.8020, 0.7819, 0.7617],
                 [0.9093, 0.9364, 0.9581, 0.9749, 0.9870],
             ]   )
         snake_case					=    TFRoFormerSinusoidalPositionalEmbedding(num_positions=5_1_2      , embedding_dim=5_1_2   )
         emba([2, 1_6, 5_1_2]   )
         snake_case					=    emba.weight[:3, :5]

         tf.debugging.assert_near(__snake_case      , __snake_case      , atol=self.tolerance   )






@require_tf
class   A__ (   unittest.TestCase     ):







    """simple docstring"""

    __magic_name__        =		1e-4
    def 		a_      (	self   ):
         # 2,12,16,64
         snake_case					=    tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4      , dtype=tf.floataa   )      , shape=(2, 1_2, 1_6, 6_4)   ) / 1_0_0

         snake_case					=    -tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4      , dtype=tf.floataa   )      , shape=(2, 1_2, 1_6, 6_4)   ) / 1_0_0

         snake_case					=    TFRoFormerSinusoidalPositionalEmbedding(num_positions=3_2      , embedding_dim=6_4   )
         snake_case					=    embed_positions([2, 1_6, 7_6_8]   )[None, None, :, :]

         snake_case			,  snake_case					=    TFRoFormerSelfAttention.apply_rotary_position_embeddings(
             __snake_case      , __snake_case      , __snake_case   )

         snake_case					=    tf.constant(
             [
                 [0.0000, 0.0100, 0.0200, 0.0300, 0.0400, 0.0500, 0.0600, 0.0700],
                 [-0.2012, 0.8897, 0.0263, 0.9401, 0.2074, 0.9463, 0.3481, 0.9343],
                 [-1.7057, 0.6271, -1.2145, 1.3897, -0.6303, 1.7647, -0.1173, 1.8985],
                 [-2.1731, -1.6397, -2.7358, 0.2854, -2.1840, 1.7183, -1.3018, 2.4871],
                 [0.2717, -3.6173, -2.9206, -2.1988, -3.6638, 0.3858, -2.9155, 2.2980],
                 [3.9859, -2.1580, -0.7984, -4.4904, -4.1181, -2.0252, -4.4782, 1.1253],
             ]   )
         snake_case					=    tf.constant(
             [
                 [0.0000, -0.0100, -0.0200, -0.0300, -0.0400, -0.0500, -0.0600, -0.0700],
                 [0.2012, -0.8897, -0.0263, -0.9401, -0.2074, -0.9463, -0.3481, -0.9343],
                 [1.7057, -0.6271, 1.2145, -1.3897, 0.6303, -1.7647, 0.1173, -1.8985],
                 [2.1731, 1.6397, 2.7358, -0.2854, 2.1840, -1.7183, 1.3018, -2.4871],
                 [-0.2717, 3.6173, 2.9206, 2.1988, 3.6638, -0.3858, 2.9155, -2.2980],
                 [-3.9859, 2.1580, 0.7984, 4.4904, 4.1181, 2.0252, 4.4782, -1.1253],
             ]   )

         tf.debugging.assert_near(query_layer[0, 0, :6, :8]      , __snake_case      , atol=self.tolerance   )
         tf.debugging.assert_near(key_layer[0, 0, :6, :8]      , __snake_case      , atol=self.tolerance   )








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_							,UpperCamelCase_=False    ):



     """simple docstring"""






     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       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_=False    ):



     """simple docstring"""






     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       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    ['''head.weight''', '''head.bias''']
     for k in ignore_keys:
          state_dict.pop(UpperCamelCase_							,UpperCamelCase_    )
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






     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_							,UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    dct.pop(UpperCamelCase_    )
     snake_case					=    val
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     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(UpperCamelCase_							,UpperCamelCase_    )							,'''r'''    )    )
     snake_case					=    {int(UpperCamelCase_    ): 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(UpperCamelCase_    )

     snake_case					=    torch.hub.load_state_dict_from_url(UpperCamelCase_							,map_location='''cpu'''    )['''target_encoder''']

     snake_case					=    ViTImageProcessor(size=config.image_size    )

     remove_projection_head(UpperCamelCase_    )
     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()

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

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

     # forward pass
     torch.manual_seed(2    )
     snake_case					=    model(**UpperCamelCase_    )
     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]							,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__":
 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   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      :			Any      				=   parser.parse_args()
 convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




