

'''simple docstring'''
from __future__ import annotations

import unittest

import numpy as np

from transformers import LayoutLMConfig, 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.models.layoutlm.modeling_tf_layoutlm import (
        TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFLayoutLMForMaskedLM,
        TFLayoutLMForQuestionAnswering,
        TFLayoutLMForSequenceClassification,
        TFLayoutLMForTokenClassification,
        TFLayoutLMModel,
    )


class    a__  :


   def __init__( self			,				_UpperCamelCase			,				_UpperCamelCase=13			,				_UpperCamelCase=7			,				_UpperCamelCase=True			,				_UpperCamelCase=True			,				_UpperCamelCase=True			,				_UpperCamelCase=True			,				_UpperCamelCase=99			,				_UpperCamelCase=32			,				_UpperCamelCase=2			,				_UpperCamelCase=4			,				_UpperCamelCase=37			,				_UpperCamelCase="gelu"			,				_UpperCamelCase=0.1			,				_UpperCamelCase=0.1			,				_UpperCamelCase=512			,				_UpperCamelCase=16			,				_UpperCamelCase=2			,				_UpperCamelCase=0.0_2			,				_UpperCamelCase=3			,				_UpperCamelCase=4			,				_UpperCamelCase=None			,				_UpperCamelCase=1000			,				):





      """simple docstring"""





      _lowercase			:      List[Any]			     =   parent
      _lowercase			:      List[str]			     =   batch_size
      _lowercase			:      List[str]			     =   seq_length
      _lowercase			:      str			     =   is_training
      _lowercase			:      List[Any]			     =   use_input_mask
      _lowercase			:      str			     =   use_token_type_ids
      _lowercase			:      Optional[Any]			     =   use_labels
      _lowercase			:      int			     =   vocab_size
      _lowercase			:      Dict			     =   hidden_size
      _lowercase			:      Tuple			     =   num_hidden_layers
      _lowercase			:      Dict			     =   num_attention_heads
      _lowercase			:      int			     =   intermediate_size
      _lowercase			:      str			     =   hidden_act
      _lowercase			:      str			     =   hidden_dropout_prob
      _lowercase			:      Any			     =   attention_probs_dropout_prob
      _lowercase			:      Optional[int]			     =   max_position_embeddings
      _lowercase			:      Union[str, Any]			     =   type_vocab_size
      _lowercase			:      Union[str, Any]			     =   type_sequence_label_size
      _lowercase			:      List[Any]			     =   initializer_range
      _lowercase			:      Dict			     =   num_labels
      _lowercase			:      Optional[Any]			     =   num_choices
      _lowercase			:      int			     =   scope
      _lowercase			:      Dict			     =   range_bbox


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Dict			     =   ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size  )

      # convert bbox to numpy since TF does not support item assignment
      _lowercase			:      List[str]			     =   ids_tensor([self.batch_size, self.seq_length, 4]			,				self.range_bbox  ).numpy()
      # Ensure that bbox is legal
      for i in range(bbox.shape[0]  ):
         for j in range(bbox.shape[1]  ):
            if bbox[i, j, 3] < bbox[i, j, 1]:
               _lowercase			:      List[str]			     =   bbox[i, j, 3]
               _lowercase			:      Dict			     =   bbox[i, j, 1]
               _lowercase			:      List[Any]			     =   t
            if bbox[i, j, 2] < bbox[i, j, 0]:
               _lowercase			:      List[Any]			     =   bbox[i, j, 2]
               _lowercase			:      List[Any]			     =   bbox[i, j, 0]
               _lowercase			:      List[Any]			     =   t
      _lowercase			:      int			     =   tf.convert_to_tensor(snake_case_  )

      _lowercase			:      int			     =   None
      if self.use_input_mask:
         _lowercase			:      Dict			     =   random_attention_mask([self.batch_size, self.seq_length]  )

      _lowercase			:      List[Any]			     =   None
      if self.use_token_type_ids:
         _lowercase			:      Optional[Any]			     =   ids_tensor([self.batch_size, self.seq_length]			,				self.type_vocab_size  )

      _lowercase			:      Optional[Any]			     =   None
      _lowercase			:      Optional[Any]			     =   None
      _lowercase			:      str			     =   None
      if self.use_labels:
         _lowercase			:      Union[str, Any]			     =   ids_tensor([self.batch_size]			,				self.type_sequence_label_size  )
         _lowercase			:      Optional[Any]			     =   ids_tensor([self.batch_size, self.seq_length]			,				self.num_labels  )
         _lowercase			:      Dict			     =   ids_tensor([self.batch_size]			,				self.num_choices  )

      _lowercase			:      List[Any]			     =   LayoutLMConfig(
          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 config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


   def 						_lowerCamelCase   ( self			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase  ):





      """simple docstring"""





      _lowercase			:      Any			     =   TFLayoutLMModel(config=snake_case_  )

      _lowercase			:      str			     =   model(snake_case_			,				snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_  )
      _lowercase			:      str			     =   model(snake_case_			,				snake_case_			,				token_type_ids=snake_case_  )
      _lowercase			:      Union[str, Any]			     =   model(snake_case_			,				snake_case_  )
      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			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase  ):





      """simple docstring"""





      _lowercase			:      List[str]			     =   TFLayoutLMForMaskedLM(config=snake_case_  )

      _lowercase			:      List[str]			     =   model(snake_case_			,				snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_			,				labels=snake_case_  )
      self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.seq_length, self.vocab_size)  )


   def 						_lowerCamelCase   ( self			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase  ):





      """simple docstring"""





      _lowercase			:      List[str]			     =   self.num_labels
      _lowercase			:      Dict			     =   TFLayoutLMForSequenceClassification(config=snake_case_  )

      _lowercase			:      int			     =   model(snake_case_			,				snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_  )
      self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.num_labels)  )


   def 						_lowerCamelCase   ( self			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase  ):





      """simple docstring"""





      _lowercase			:      List[str]			     =   self.num_labels
      _lowercase			:      List[Any]			     =   TFLayoutLMForTokenClassification(config=snake_case_  )

      _lowercase			:      Any			     =   model(snake_case_			,				snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_			,				labels=snake_case_  )
      self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.seq_length, self.num_labels)  )


   def 						_lowerCamelCase   ( self			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase			,				_UpperCamelCase  ):





      """simple docstring"""





      _lowercase			:      Union[str, Any]			     =   TFLayoutLMForQuestionAnswering(config=snake_case_  )

      _lowercase			:      Tuple			     =   model(snake_case_			,				snake_case_			,				attention_mask=snake_case_			,				token_type_ids=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 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Dict			     =   self.prepare_config_and_inputs()
      (
          _lowercase       
      )			:      int			     =   config_and_inputs
      _lowercase			:      Optional[int]			     =   {
          '''input_ids''': input_ids,
          '''bbox''': bbox,
          '''token_type_ids''': token_type_ids,
          '''attention_mask''': input_mask,
      }
      return config, inputs_dict




@require_tf
class    a__  (							__lowercase  ,      __lowercase  ,      unittest.TestCase  ):
   _SCREAMING_SNAKE_CASE    :     Tuple	       =							(
       (
           TFLayoutLMModel,
           TFLayoutLMForMaskedLM,
           TFLayoutLMForTokenClassification,
           TFLayoutLMForSequenceClassification,
           TFLayoutLMForQuestionAnswering,
       )
       if is_tf_available()
       else ()
   )
   _SCREAMING_SNAKE_CASE    :     List[str]	       =							(
       {
           "feature-extraction": TFLayoutLMModel,
           "fill-mask": TFLayoutLMForMaskedLM,
           "text-classification": TFLayoutLMForSequenceClassification,
           "token-classification": TFLayoutLMForTokenClassification,
           "zero-shot": TFLayoutLMForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   _SCREAMING_SNAKE_CASE    :     List[Any]	       =							False
   _SCREAMING_SNAKE_CASE    :     Dict	       =							True
   _SCREAMING_SNAKE_CASE    :     Optional[int]	       =							10


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Union[str, Any]			     =   TFLayoutLMModelTester(self  )
      _lowercase			:      Optional[Any]			     =   ConfigTester(self			,				config_class=snake_case_			,				hidden_size=37  )


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      self.config_tester.run_common_tests()


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      List[str]			     =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*snake_case_  )


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      int			     =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*snake_case_  )


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Dict			     =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*snake_case_  )


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Tuple			     =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*snake_case_  )


   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      List[str]			     =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*snake_case_  )


   @slow
   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      for model_name in TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _lowercase			:      Optional[Any]			     =   TFLayoutLMModel.from_pretrained(snake_case_  )
         self.assertIsNotNone(snake_case_  )






   @unittest.skip("Onnx compliancy broke with TF 2.10"  )
   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      pass







def 					_A		(	)			->	str:
   # Here we prepare a batch of 2 sequences to test a LayoutLM forward pass on:
   # fmt: off
   _lowercase			:      List[str]			     =   tf.convert_to_tensor([[1_01,10_19,10_14,10_16,10_37,1_28_49,47_47,10_04,1_42_46,22_78,54_39,45_24,50_02,29_30,21_93,29_30,43_41,32_08,10_05,10_55,21_71,28_48,1_13_00,35_31,1_02],[1_01,40_70,40_34,70_20,10_24,30_58,10_15,10_13,28_61,10_13,60_70,1_92_74,27_72,62_05,2_78_14,1_61_47,1_61_47,43_43,20_47,1_02_83,1_09_69,1_43_89,10_12,23_38,1_02]]     )  # noqa: E231
   _lowercase			:      Union[str, Any]			     =   tf.convert_to_tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]     )  # noqa: E231
   _lowercase			:      str			     =   tf.convert_to_tensor([[[0,0,0,0],[4_23,2_37,4_40,2_51],[4_27,2_72,4_41,2_87],[4_19,1_15,4_37,1_29],[9_61,8_85,9_92,9_12],[2_56,38,3_30,58],[2_56,38,3_30,58],[3_36,42,3_53,57],[3_60,39,4_01,56],[3_60,39,4_01,56],[4_11,39,4_71,59],[4_79,41,5_28,59],[5_33,39,6_30,60],[67,1_13,1_34,1_31],[1_41,1_15,2_09,1_32],[68,1_49,1_33,1_66],[1_41,1_49,1_87,1_64],[1_95,1_48,2_87,1_65],[1_95,1_48,2_87,1_65],[1_95,1_48,2_87,1_65],[2_95,1_48,3_49,1_65],[4_41,1_49,4_92,1_66],[4_97,1_49,5_46,1_64],[64,2_01,1_25,2_18],[10_00,10_00,10_00,10_00]],[[0,0,0,0],[6_62,1_50,7_54,1_66],[6_65,1_99,7_42,2_11],[5_19,2_13,5_54,2_28],[5_19,2_13,5_54,2_28],[1_34,4_33,1_87,4_54],[1_30,4_67,2_04,4_80],[1_30,4_67,2_04,4_80],[1_30,4_67,2_04,4_80],[1_30,4_67,2_04,4_80],[1_30,4_67,2_04,4_80],[3_14,4_69,3_76,4_82],[5_04,6_84,5_82,7_06],[9_41,8_25,9_73,9_00],[9_41,8_25,9_73,9_00],[9_41,8_25,9_73,9_00],[9_41,8_25,9_73,9_00],[6_10,7_49,6_52,7_65],[1_30,6_59,1_68,6_72],[1_76,6_57,2_37,6_72],[2_38,6_57,3_12,6_72],[4_43,6_53,6_28,6_72],[4_43,6_53,6_28,6_72],[7_16,3_01,8_25,3_17],[10_00,10_00,10_00,10_00]]]     )  # noqa: E231
   _lowercase			:      Any			     =   tf.convert_to_tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]     )  # noqa: E231
   # these are sequence labels (i.e. at the token level)
   _lowercase			:      str			     =   tf.convert_to_tensor([[-1_00,10,10,10,9,1,-1_00,7,7,-1_00,7,7,4,2,5,2,8,8,-1_00,-1_00,5,0,3,2,-1_00],[-1_00,12,12,12,-1_00,12,10,-1_00,-1_00,-1_00,-1_00,10,12,9,-1_00,-1_00,-1_00,10,10,10,9,12,-1_00,10,-1_00]]     )  # noqa: E231
   # fmt: on

   return input_ids, attention_mask, bbox, token_type_ids, labels


@require_tf
class    a__  (							unittest.TestCase  ):


   @slow
   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Tuple			     =   TFLayoutLMModel.from_pretrained("microsoft/layoutlm-base-uncased"  )

      _lowercase			:      Any			     =   prepare_layoutlm_batch_inputs()

      # forward pass
      _lowercase			:      int			     =   model(input_ids=snake_case_			,				bbox=snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_  )

      # test the sequence output on [0, :3, :3]
      _lowercase			:      Optional[Any]			     =   tf.convert_to_tensor(
          [[0.1_7_8_5, -0.1_9_4_7, -0.0_4_2_5], [-0.3_2_5_4, -0.2_8_0_7, 0.2_5_5_3], [-0.5_3_9_1, -0.3_3_2_2, 0.3_3_6_4]]			,				)

      self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]			,				snake_case_			,				atol=1E-3  )  )

      # test the pooled output on [1, :3]
      _lowercase			:      int			     =   tf.convert_to_tensor([-0.6_5_8_0, -0.0_2_1_4, 0.8_5_5_2]  )

      self.assertTrue(np.allclose(outputs.pooler_output[1, :3]			,				snake_case_			,				atol=1E-3  )  )


   @slow
   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      int			     =   TFLayoutLMForSequenceClassification.from_pretrained("microsoft/layoutlm-base-uncased"			,				num_labels=2  )

      _lowercase			:      str			     =   prepare_layoutlm_batch_inputs()

      # forward pass
      _lowercase			:      Union[str, Any]			     =   model(
          input_ids=snake_case_			,				bbox=snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_			,				labels=tf.convert_to_tensor([1, 1]  )			,				)

      # test whether we get a loss as a scalar
      _lowercase			:      Tuple			     =   outputs.loss
      _lowercase			:      Dict			     =   (2,)
      self.assertEqual(loss.shape			,				snake_case_  )

      # test the shape of the logits
      _lowercase			:      Dict			     =   outputs.logits
      _lowercase			:      List[str]			     =   (2, 2)
      self.assertEqual(logits.shape			,				snake_case_  )


   @slow
   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Union[str, Any]			     =   TFLayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased"			,				num_labels=13  )

      _lowercase			:      Optional[int]			     =   prepare_layoutlm_batch_inputs()

      # forward pass
      _lowercase			:      Optional[int]			     =   model(
          input_ids=snake_case_			,				bbox=snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_			,				labels=snake_case_  )

      # test the shape of the logits
      _lowercase			:      str			     =   outputs.logits
      _lowercase			:      List[str]			     =   tf.convert_to_tensor((2, 25, 13)  )
      self.assertEqual(logits.shape			,				snake_case_  )






   @slow
   def 						_lowerCamelCase   ( self  ):





      """simple docstring"""





      _lowercase			:      Union[str, Any]			     =   TFLayoutLMForQuestionAnswering.from_pretrained("microsoft/layoutlm-base-uncased"  )

      _lowercase			:      Dict			     =   prepare_layoutlm_batch_inputs()

      # forward pass
      _lowercase			:      str			     =   model(input_ids=snake_case_			,				bbox=snake_case_			,				attention_mask=snake_case_			,				token_type_ids=snake_case_  )

      # test the shape of the logits
      _lowercase			:      Optional[Any]			     =   tf.convert_to_tensor((2, 25)  )
      self.assertEqual(outputs.start_logits.shape			,				snake_case_  )
      self.assertEqual(outputs.end_logits.shape			,				snake_case_  )








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     (				lowerCAmelCase__							:				Optional[int]  ,   lowerCAmelCase__							:				str=False  ):
 __a :  Any						=  []
 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"
  __a :  Optional[int]						=  [(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     (				lowerCAmelCase__							:				Dict  ,   lowerCAmelCase__							:				Union[str, Any]  ,   lowerCAmelCase__							:				List[Any]=False  ):
 for i in range(config.num_hidden_layers  ):
  if base_model:
   __a :  Dict						=  ''''''
  else:
   __a :  str						=  '''vit.'''
  # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
  __a :  List[Any]						=  state_dict.pop(f"module.blocks.{i}.attn.qkv.weight"  )
  __a :  Tuple						=  state_dict.pop(f"module.blocks.{i}.attn.qkv.bias"  )
  # next, add query, keys and values (in that order) to the state dict
  __a :  Dict						=  in_proj_weight[
      : config.hidden_size, :
  ]
  __a :  str						=  in_proj_bias[: config.hidden_size]
  __a :  int						=  in_proj_weight[
      config.hidden_size : config.hidden_size * 2, :
  ]
  __a :  Dict						=  in_proj_bias[
      config.hidden_size : config.hidden_size * 2
  ]
  __a :  str						=  in_proj_weight[
      -config.hidden_size :, :
  ]
  __a :  Optional[Any]						=  in_proj_bias[-config.hidden_size :]




def    __UpperCamelCase     (				lowerCAmelCase__							:				Optional[int]  ):
 __a :  Dict						=  ['''head.weight''', '''head.bias''']
 for k in ignore_keys:
  state_dict.pop(lowerCAmelCase__  ,   lowerCAmelCase__  )




def    __UpperCamelCase     (				lowerCAmelCase__							:				Optional[Any]  ):
 # projection head is used in the self-supervised pre-training in MSN,
 # for downstream task it's not needed.
 __a :  str						=  [
     '''module.fc.fc1.weight''',
     '''module.fc.fc1.bias''',
     '''module.fc.bn1.weight''',
     '''module.fc.bn1.bias''',
     '''module.fc.bn1.running_mean''',
     '''module.fc.bn1.running_var''',
     '''module.fc.bn1.num_batches_tracked''',
     '''module.fc.fc2.weight''',
     '''module.fc.fc2.bias''',
     '''module.fc.bn2.weight''',
     '''module.fc.bn2.bias''',
     '''module.fc.bn2.running_mean''',
     '''module.fc.bn2.running_var''',
     '''module.fc.bn2.num_batches_tracked''',
     '''module.fc.fc3.weight''',
     '''module.fc.fc3.bias''',
 ]
 for k in ignore_keys:
  state_dict.pop(lowerCAmelCase__  ,   lowerCAmelCase__  )




def    __UpperCamelCase     (				lowerCAmelCase__							:				Dict  ,   lowerCAmelCase__							:				List[Any]  ,   lowerCAmelCase__							:				List[Any]  ):
 __a :  Optional[Any]						=  dct.pop(lowerCAmelCase__  )
 __a :  Optional[Any]						=  val




def    __UpperCamelCase     (				lowerCAmelCase__							:				Optional[Any]  ,   lowerCAmelCase__							:				str  ):
 __a :  str						=  ViTMSNConfig()
 __a :  List[Any]						=  1_0_0_0

 __a :  Union[str, Any]						=  '''datasets/huggingface/label-files'''
 __a :  Optional[int]						=  '''imagenet-1k-id2label.json'''
 __a :  Optional[int]						=  json.load(open(hf_hub_download(lowerCAmelCase__  ,   lowerCAmelCase__  )  ,   '''r'''  )  )
 __a :  Dict						=  {int(lowerCAmelCase__  ): v for k, v in idalabel.items()}
 __a :  Tuple						=  idalabel
 __a :  int						=  {v: k for k, v in idalabel.items()}

 if "s16" in checkpoint_url:
  __a :  Union[str, Any]						=  3_8_4
  __a :  Union[str, Any]						=  1_5_3_6
  __a :  Union[str, Any]						=  6
 elif "l16" in checkpoint_url:
  __a :  str						=  1_0_2_4
  __a :  Union[str, Any]						=  4_0_9_6
  __a :  Optional[int]						=  2_4
  __a :  int						=  1_6
  __a :  List[Any]						=  0.1
 elif "b4" in checkpoint_url:
  __a :  int						=  4
 elif "l7" in checkpoint_url:
  __a :  int						=  7
  __a :  List[Any]						=  1_0_2_4
  __a :  Union[str, Any]						=  4_0_9_6
  __a :  List[str]						=  2_4
  __a :  int						=  1_6
  __a :  Tuple						=  0.1

 __a :  Dict						=  ViTMSNModel(lowerCAmelCase__  )

 __a :  Union[str, Any]						=  torch.hub.load_state_dict_from_url(lowerCAmelCase__  ,   map_location='''cpu'''  )['''target_encoder''']

 __a :  Optional[int]						=  ViTImageProcessor(size=config.image_size  )

 remove_projection_head(lowerCAmelCase__  )
 __a :  str						=  create_rename_keys(lowerCAmelCase__  ,   base_model=lowerCAmelCase__  )

 for src, dest in rename_keys:
  rename_key(lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__  )
 read_in_q_k_v(lowerCAmelCase__  ,   lowerCAmelCase__  ,   base_model=lowerCAmelCase__  )

 model.load_state_dict(lowerCAmelCase__  )
 model.eval()

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

 __a :  str						=  Image.open(requests.get(lowerCAmelCase__  ,   stream=lowerCAmelCase__  ).raw  )
 __a :  Any						=  ViTImageProcessor(
     size=config.image_size  ,   image_mean=lowerCAmelCase__  ,   image_std=lowerCAmelCase__  )
 __a :  Tuple						=  image_processor(images=lowerCAmelCase__  ,   return_tensors='''pt'''  )

 # forward pass
 torch.manual_seed(2  )
 __a :  str						=  model(**lowerCAmelCase__  )
 __a :  List[Any]						=  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:
  __a :  Dict						=  torch.tensor([[-1.09_15, -1.48_76, -1.18_09]]  )
 elif "b16" in checkpoint_url:
  __a :  List[str]						=  torch.tensor([[14.28_89, -18.90_45, 11.72_81]]  )
 elif "l16" in checkpoint_url:
  __a :  str						=  torch.tensor([[41.50_28, -22.86_81, 45.64_75]]  )
 elif "b4" in checkpoint_url:
  __a :  Union[str, Any]						=  torch.tensor([[-4.38_68, 5.29_32, -0.41_37]]  )
 else:
  __a :  Dict						=  torch.tensor([[-0.17_92, -0.64_65, 2.42_63]]  )

 # verify logits
 assert torch.allclose(last_hidden_state[:, 0, :3]  ,   lowerCAmelCase__  ,   atol=1e-4  )

 print(f"Saving model to {pytorch_dump_folder_path}"  )
 model.save_pretrained(lowerCAmelCase__  )

 print(f"Saving image processor to {pytorch_dump_folder_path}"  )
 image_processor.save_pretrained(lowerCAmelCase__  )


if __name__ == "__main__":
      lowercase__       =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.'
      )

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