





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 							__SCREAMING_SNAKE_CASE       :






							def __init__(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=13	,			__lowerCAmelCase=7	,			__lowerCAmelCase=True	,			__lowerCAmelCase=True	,			__lowerCAmelCase=True	,			__lowerCAmelCase=True	,			__lowerCAmelCase=99	,			__lowerCAmelCase=32	,			__lowerCAmelCase=2	,			__lowerCAmelCase=4	,			__lowerCAmelCase=37	,			__lowerCAmelCase="gelu"	,			__lowerCAmelCase=0.1	,			__lowerCAmelCase=0.1	,			__lowerCAmelCase=512	,			__lowerCAmelCase=16	,			__lowerCAmelCase=2	,			__lowerCAmelCase=0.02	,			__lowerCAmelCase=3	,			__lowerCAmelCase=4	,			__lowerCAmelCase=None	,			__lowerCAmelCase=1000	,			):
														UpperCamelCase__					  =      parent
														UpperCamelCase__					  =      batch_size
														UpperCamelCase__					  =      seq_length
														UpperCamelCase__					  =      is_training
														UpperCamelCase__					  =      use_input_mask
														UpperCamelCase__					  =      use_token_type_ids
														UpperCamelCase__					  =      use_labels
														UpperCamelCase__					  =      vocab_size
														UpperCamelCase__					  =      hidden_size
														UpperCamelCase__					  =      num_hidden_layers
														UpperCamelCase__					  =      num_attention_heads
														UpperCamelCase__					  =      intermediate_size
														UpperCamelCase__					  =      hidden_act
														UpperCamelCase__					  =      hidden_dropout_prob
														UpperCamelCase__					  =      attention_probs_dropout_prob
														UpperCamelCase__					  =      max_position_embeddings
														UpperCamelCase__					  =      type_vocab_size
														UpperCamelCase__					  =      type_sequence_label_size
														UpperCamelCase__					  =      initializer_range
														UpperCamelCase__					  =      num_labels
														UpperCamelCase__					  =      num_choices
														UpperCamelCase__					  =      scope
														UpperCamelCase__					  =      range_bbox






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      ids_tensor([self.batch_size, self.seq_length]	,			self.vocab_size	)

														# convert bbox to numpy since TF does not support item assignment
														UpperCamelCase__					  =      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]:
																																			UpperCamelCase__					  =      bbox[i, j, 3]
																																			UpperCamelCase__					  =      bbox[i, j, 1]
																																			UpperCamelCase__					  =      t
																												if bbox[i, j, 2] < bbox[i, j, 0]:
																																			UpperCamelCase__					  =      bbox[i, j, 2]
																																			UpperCamelCase__					  =      bbox[i, j, 0]
																																			UpperCamelCase__					  =      t
														UpperCamelCase__					  =      tf.convert_to_tensor(__lowerCAmelCase	)

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

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

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

														UpperCamelCase__					  =      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	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      TFLayoutLMModel(config=__lowerCAmelCase	)

														UpperCamelCase__					  =      model(__lowerCAmelCase	,			__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	)
														UpperCamelCase__					  =      model(__lowerCAmelCase	,			__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	)
														UpperCamelCase__					  =      model(__lowerCAmelCase	,			__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	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      TFLayoutLMForMaskedLM(config=__lowerCAmelCase	)

														UpperCamelCase__					  =      model(__lowerCAmelCase	,			__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	,			labels=__lowerCAmelCase	)
														self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.seq_length, self.vocab_size)	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      self.num_labels
														UpperCamelCase__					  =      TFLayoutLMForSequenceClassification(config=__lowerCAmelCase	)

														UpperCamelCase__					  =      model(__lowerCAmelCase	,			__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	)
														self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.num_labels)	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      self.num_labels
														UpperCamelCase__					  =      TFLayoutLMForTokenClassification(config=__lowerCAmelCase	)

														UpperCamelCase__					  =      model(__lowerCAmelCase	,			__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	,			labels=__lowerCAmelCase	)
														self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.seq_length, self.num_labels)	)






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														UpperCamelCase__					  =      TFLayoutLMForQuestionAnswering(config=__lowerCAmelCase	)

														UpperCamelCase__					  =      model(__lowerCAmelCase	,			__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__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	):
														UpperCamelCase__					  =      self.prepare_config_and_inputs()
														(
														    (
														    UpperCamelCase__
														)  ,  (
														    UpperCamelCase__
														)  ,  (
														    UpperCamelCase__
														)  ,  (
														    UpperCamelCase__
														)  ,  (
														    UpperCamelCase__
														)  ,  (
														    UpperCamelCase__
														)  ,  (
														    UpperCamelCase__
														)  ,  (
														    UpperCamelCase__
														)  ,  
														)					  =      config_and_inputs
														UpperCamelCase__					  =      {
														    """input_ids""": input_ids,
														    """bbox""": bbox,
														    """token_type_ids""": token_type_ids,
														    """attention_mask""": input_mask,
														}
														return config, inputs_dict






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






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      TFLayoutLMModelTester(self	)
														UpperCamelCase__					  =      ConfigTester(self	,			config_class=__lowerCAmelCase	,			hidden_size=37	)






							def 					_lowerCamelCase		(    self	):
														self.config_tester.run_common_tests()






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase	)






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase	)






							@slow
							def 					_lowerCamelCase		(    self	):
														for model_name in TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					UpperCamelCase__					  =      TFLayoutLMModel.from_pretrained(__lowerCAmelCase	)
																					self.assertIsNotNone(__lowerCAmelCase	)







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






def      _UpperCamelCase  ():






							"""simple docstring"""





							UpperCamelCase__					  =      tf.convert_to_tensor([[101,1019,1014,1016,1037,1_2849,4747,1004,1_4246,2278,5439,4524,5002,2930,2193,2930,4341,3208,1005,1055,2171,2848,1_1300,3531,102],[101,4070,4034,7020,1024,3058,1015,1013,2861,1013,6070,1_9274,2772,6205,2_7814,1_6147,1_6147,4343,2047,1_0283,1_0969,1_4389,1012,2338,102]]							)  # noqa: E231
							UpperCamelCase__					  =      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
							UpperCamelCase__					  =      tf.convert_to_tensor([[[0,0,0,0],[423,237,440,251],[427,272,441,287],[419,115,437,129],[961,885,992,912],[256,38,330,58],[256,38,330,58],[336,42,353,57],[360,39,401,56],[360,39,401,56],[411,39,471,59],[479,41,528,59],[533,39,630,60],[67,113,134,131],[141,115,209,132],[68,149,133,166],[141,149,187,164],[195,148,287,165],[195,148,287,165],[195,148,287,165],[295,148,349,165],[441,149,492,166],[497,149,546,164],[64,201,125,218],[1000,1000,1000,1000]],[[0,0,0,0],[662,150,754,166],[665,199,742,211],[519,213,554,228],[519,213,554,228],[134,433,187,454],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[314,469,376,482],[504,684,582,706],[941,825,973,900],[941,825,973,900],[941,825,973,900],[941,825,973,900],[610,749,652,765],[130,659,168,672],[176,657,237,672],[238,657,312,672],[443,653,628,672],[443,653,628,672],[716,301,825,317],[1000,1000,1000,1000]]]							)  # noqa: E231
							UpperCamelCase__					  =      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)
							UpperCamelCase__					  =      tf.convert_to_tensor([[-100,10,10,10,9,1,-100,7,7,-100,7,7,4,2,5,2,8,8,-100,-100,5,0,3,2,-100],[-100,12,12,12,-100,12,10,-100,-100,-100,-100,10,12,9,-100,-100,-100,10,10,10,9,12,-100,10,-100]]							)  # noqa: E231
							# fmt: on

							return input_ids, attention_mask, bbox, token_type_ids, labels

@require_tf
class 							__SCREAMING_SNAKE_CASE       (  unittest.TestCase							):






							@slow
							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      TFLayoutLMModel.from_pretrained("""microsoft/layoutlm-base-uncased"""	)

														UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__					  =      prepare_layoutlm_batch_inputs()

														# forward pass
														UpperCamelCase__					  =      model(input_ids=__lowerCAmelCase	,			bbox=__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	)

														# test the sequence output on [0, :3, :3]
														UpperCamelCase__					  =      tf.convert_to_tensor(
														    [[0.1785, -0.1947, -0.0425], [-0.3254, -0.2807, 0.2553], [-0.5391, -0.3322, 0.3364]]	,			)

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

														# test the pooled output on [1, :3]
														UpperCamelCase__					  =      tf.convert_to_tensor([-0.6580, -0.0214, 0.8552]	)

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






							@slow
							def 					_lowerCamelCase		(    self	):
														# initialize model with randomly initialized sequence classification head
														UpperCamelCase__					  =      TFLayoutLMForSequenceClassification.from_pretrained("""microsoft/layoutlm-base-uncased"""	,			num_labels=2	)

														UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__					  =      prepare_layoutlm_batch_inputs()

														# forward pass
														UpperCamelCase__					  =      model(
														    input_ids=__lowerCAmelCase	,			bbox=__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	,			labels=tf.convert_to_tensor([1, 1]	)	,			)

														# test whether we get a loss as a scalar
														UpperCamelCase__					  =      outputs.loss
														UpperCamelCase__					  =      (2,)
														self.assertEqual(loss.shape	,			__lowerCAmelCase	)

														# test the shape of the logits
														UpperCamelCase__					  =      outputs.logits
														UpperCamelCase__					  =      (2, 2)
														self.assertEqual(logits.shape	,			__lowerCAmelCase	)






							@slow
							def 					_lowerCamelCase		(    self	):
														# initialize model with randomly initialized token classification head
														UpperCamelCase__					  =      TFLayoutLMForTokenClassification.from_pretrained("""microsoft/layoutlm-base-uncased"""	,			num_labels=13	)

														UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__					  =      prepare_layoutlm_batch_inputs()

														# forward pass
														UpperCamelCase__					  =      model(
														    input_ids=__lowerCAmelCase	,			bbox=__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	,			labels=__lowerCAmelCase	)

														# test the shape of the logits
														UpperCamelCase__					  =      outputs.logits
														UpperCamelCase__					  =      tf.convert_to_tensor((2, 25, 13)	)
														self.assertEqual(logits.shape	,			__lowerCAmelCase	)







							@slow
							def 					_lowerCamelCase		(    self	):
														# initialize model with randomly initialized token classification head
														UpperCamelCase__					  =      TFLayoutLMForQuestionAnswering.from_pretrained("""microsoft/layoutlm-base-uncased"""	)

														UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__  ,  UpperCamelCase__					  =      prepare_layoutlm_batch_inputs()

														# forward pass
														UpperCamelCase__					  =      model(input_ids=__lowerCAmelCase	,			bbox=__lowerCAmelCase	,			attention_mask=__lowerCAmelCase	,			token_type_ids=__lowerCAmelCase	)

														# test the shape of the logits
														UpperCamelCase__					  =      tf.convert_to_tensor((2, 25)	)
														self.assertEqual(outputs.start_logits.shape	,			__lowerCAmelCase	)
														self.assertEqual(outputs.end_logits.shape	,			__lowerCAmelCase	)











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  (a__ :Tuple ,      a__ :List[str]=False							):






							"""simple docstring"""





							UpperCamelCase__					  =      []
							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"
														UpperCamelCase__					  =      [(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  (a__ :int ,      a__ :Tuple ,      a__ :Union[str, Any]=False							):






							"""simple docstring"""





							for i in range(config.num_hidden_layers							):
														if base_model:
																					UpperCamelCase__					  =      """"""
														else:
																					UpperCamelCase__					  =      """vit."""
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														UpperCamelCase__					  =      state_dict.pop(f"""module.blocks.{i}.attn.qkv.weight"""							)
														UpperCamelCase__					  =      state_dict.pop(f"""module.blocks.{i}.attn.qkv.bias"""							)
														# next, add query, keys and values (in that order) to the state dict
														UpperCamelCase__					  =      in_proj_weight[
														    : config.hidden_size, :
														]
														UpperCamelCase__					  =      in_proj_bias[: config.hidden_size]
														UpperCamelCase__					  =      in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														UpperCamelCase__					  =      in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														UpperCamelCase__					  =      in_proj_weight[
														    -config.hidden_size :, :
														]
														UpperCamelCase__					  =      in_proj_bias[-config.hidden_size :]





def      _UpperCamelCase  (a__ :int							):






							"""simple docstring"""





							UpperCamelCase__					  =      ["""head.weight""", """head.bias"""]
							for k in ignore_keys:
														state_dict.pop(a__ ,      a__							)





def      _UpperCamelCase  (a__ :Optional[int]							):






							"""simple docstring"""





							UpperCamelCase__					  =      [
							    """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(a__ ,      a__							)





def      _UpperCamelCase  (a__ :List[str] ,      a__ :Any ,      a__ :str							):






							"""simple docstring"""





							UpperCamelCase__					  =      dct.pop(a__							)
							UpperCamelCase__					  =      val





def      _UpperCamelCase  (a__ :Dict ,      a__ :Dict							):






							"""simple docstring"""





							UpperCamelCase__					  =      ViTMSNConfig()
							UpperCamelCase__					  =      1000

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

							if "s16" in checkpoint_url:
														UpperCamelCase__					  =      384
														UpperCamelCase__					  =      1536
														UpperCamelCase__					  =      6
							elif "l16" in checkpoint_url:
														UpperCamelCase__					  =      1024
														UpperCamelCase__					  =      4096
														UpperCamelCase__					  =      24
														UpperCamelCase__					  =      16
														UpperCamelCase__					  =      0.1
							elif "b4" in checkpoint_url:
														UpperCamelCase__					  =      4
							elif "l7" in checkpoint_url:
														UpperCamelCase__					  =      7
														UpperCamelCase__					  =      1024
														UpperCamelCase__					  =      4096
														UpperCamelCase__					  =      24
														UpperCamelCase__					  =      16
														UpperCamelCase__					  =      0.1

							UpperCamelCase__					  =      ViTMSNModel(a__							)

							UpperCamelCase__					  =      torch.hub.load_state_dict_from_url(a__ ,      map_location="""cpu"""							)["""target_encoder"""]

							UpperCamelCase__					  =      ViTImageProcessor(size=config.image_size							)

							remove_projection_head(a__							)
							UpperCamelCase__					  =      create_rename_keys(a__ ,      base_model=a__							)

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

							model.load_state_dict(a__							)
							model.eval()

							UpperCamelCase__					  =      """http://images.cocodataset.org/val2017/000000039769.jpg"""

							UpperCamelCase__					  =      Image.open(requests.get(a__ ,      stream=a__							).raw							)
							UpperCamelCase__					  =      ViTImageProcessor(
							    size=config.image_size ,      image_mean=a__ ,      image_std=a__							)
							UpperCamelCase__					  =      image_processor(images=a__ ,      return_tensors="""pt"""							)

							# forward pass
							torch.manual_seed(2							)
							UpperCamelCase__					  =      model(**a__							)
							UpperCamelCase__					  =      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:
														UpperCamelCase__					  =      torch.tensor([[-1.0915, -1.4876, -1.1809]]							)
							elif "b16" in checkpoint_url:
														UpperCamelCase__					  =      torch.tensor([[14.2889, -18.9045, 11.7281]]							)
							elif "l16" in checkpoint_url:
														UpperCamelCase__					  =      torch.tensor([[41.5028, -22.8681, 45.6475]]							)
							elif "b4" in checkpoint_url:
														UpperCamelCase__					  =      torch.tensor([[-4.3868, 5.2932, -0.4137]]							)
							else:
														UpperCamelCase__					  =      torch.tensor([[-0.1792, -0.6465, 2.4263]]							)

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

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

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


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

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




