
import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
	import torch

	from transformers import (
	    MODEL_FOR_PRETRAINING_MAPPING,
	    FunnelBaseModel,
	    FunnelForMaskedLM,
	    FunnelForMultipleChoice,
	    FunnelForPreTraining,
	    FunnelForQuestionAnswering,
	    FunnelForSequenceClassification,
	    FunnelForTokenClassification,
	    FunnelModel,
	)







class 	A_       :




	'''simple docstring'''


	def __init__(	self		,   snake_case		,   snake_case=13		,   snake_case=7		,   snake_case=True		,   snake_case=True		,   snake_case=True		,   snake_case=True		,   snake_case=99		,   snake_case=[1, 1, 2]		,   snake_case=1		,   snake_case=32		,   snake_case=4		,   snake_case=8		,   snake_case=37		,   snake_case="gelu_new"		,   snake_case=0.1		,   snake_case=0.1		,   snake_case=0.0		,   snake_case=512		,   snake_case=3		,   snake_case=0.02		,   snake_case=3		,   snake_case=4		,   snake_case=None		,   snake_case=False		,   ):
		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					  =   block_sizes
		lowercase					  =   num_decoder_layers
		lowercase					  =   d_model
		lowercase					  =   n_head
		lowercase					  =   d_head
		lowercase					  =   d_inner
		lowercase					  =   hidden_act
		lowercase					  =   hidden_dropout
		lowercase					  =   attention_dropout
		lowercase					  =   activation_dropout
		lowercase					  =   max_position_embeddings
		lowercase					  =   type_vocab_size
		lowercase					  =   2
		lowercase					  =   num_labels
		lowercase					  =   num_choices
		lowercase					  =   scope
		lowercase					  =   initializer_std

		# Used in the tests to check the size of the first attention layer
		lowercase					  =   n_head
		# Used in the tests to check the size of the first hidden state
		lowercase					  =   self.d_model
		# Used in the tests to check the number of output hidden states/attentions
		lowercase					  =   sum(self.block_sizes	) + (0 if base else self.num_decoder_layers)
		# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
		# the last hidden state of the first block (which is the first hidden state of the decoder).
		if not base:
			lowercase					  =   self.num_hidden_layers + 2





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   ids_tensor([self.batch_size, self.seq_length]		,   self.vocab_size	)

		lowercase					  =   None
		if self.use_input_mask:
			lowercase					  =   ids_tensor([self.batch_size, self.seq_length]		,   vocab_size=2	)

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

		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					  =   ids_tensor([self.batch_size, self.seq_length]		,   1	)

		lowercase					  =   self.get_config()

		return (
		    config,
		    input_ids,
		    token_type_ids,
		    input_mask,
		    sequence_labels,
		    token_labels,
		    choice_labels,
		    fake_token_labels,
		)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		return FunnelConfig(
		    vocab_size=self.vocab_size		,   block_sizes=self.block_sizes		,   num_decoder_layers=self.num_decoder_layers		,   d_model=self.d_model		,   n_head=self.n_head		,   d_head=self.d_head		,   d_inner=self.d_inner		,   hidden_act=self.hidden_act		,   hidden_dropout=self.hidden_dropout		,   attention_dropout=self.attention_dropout		,   activation_dropout=self.activation_dropout		,   max_position_embeddings=self.max_position_embeddings		,   type_vocab_size=self.type_vocab_size		,   initializer_std=self.initializer_std		,   )





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   FunnelModel(config=snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   model(snake_case		,   attention_mask=snake_case		,   token_type_ids=snake_case	)
		lowercase					  =   model(snake_case		,   token_type_ids=snake_case	)
		lowercase					  =   model(snake_case	)
		self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, self.seq_length, self.d_model)	)

		lowercase					  =   False
		lowercase					  =   model(snake_case	)
		self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, self.seq_length, self.d_model)	)

		lowercase					  =   False
		lowercase					  =   model(snake_case	)
		self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, self.seq_length, self.d_model)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   FunnelBaseModel(config=snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   model(snake_case		,   attention_mask=snake_case		,   token_type_ids=snake_case	)
		lowercase					  =   model(snake_case		,   token_type_ids=snake_case	)
		lowercase					  =   model(snake_case	)
		self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, 2, self.d_model)	)

		lowercase					  =   False
		lowercase					  =   model(snake_case	)
		self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, 3, self.d_model)	)

		lowercase					  =   False
		lowercase					  =   model(snake_case	)
		self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, 2, self.d_model)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   self.num_labels
		lowercase					  =   FunnelForPreTraining(config=snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   model(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)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   FunnelForMaskedLM(config=snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   model(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      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   self.num_labels
		lowercase					  =   FunnelForSequenceClassification(snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   model(snake_case		,   attention_mask=snake_case		,   token_type_ids=snake_case		,   labels=snake_case	)
		self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.num_labels)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   self.num_choices
		lowercase					  =   FunnelForMultipleChoice(config=snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   input_ids.unsqueeze(1	).expand(-1		,   self.num_choices		,   -1	).contiguous()
		lowercase					  =   token_type_ids.unsqueeze(1	).expand(-1		,   self.num_choices		,   -1	).contiguous()
		lowercase					  =   input_mask.unsqueeze(1	).expand(-1		,   self.num_choices		,   -1	).contiguous()
		lowercase					  =   model(
		    snake_case		,   attention_mask=snake_case		,   token_type_ids=snake_case		,   labels=snake_case		,   )
		self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.num_choices)	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   self.num_labels
		lowercase					  =   FunnelForTokenClassification(config=snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   model(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      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   snake_case		,   ):
		lowercase					  =   FunnelForQuestionAnswering(config=snake_case	)
		model.to(snake_case	)
		model.eval()
		lowercase					  =   model(
		    snake_case		,   attention_mask=snake_case		,   token_type_ids=snake_case		,   start_positions=snake_case		,   end_positions=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      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.prepare_config_and_inputs()
		(
		    lowercase						
		)					  =   config_and_inputs
		lowercase					  =   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
		return config, inputs_dict







@require_torch
class 	A_       (      UpperCAmelCase__				,     UpperCAmelCase__				,     unittest.TestCase       ):




	'''simple docstring'''


	_UpperCamelCase					:					Any										=	False
	_UpperCamelCase					:					str										=	False
	_UpperCamelCase					:					Tuple										=	(
	    (
	        FunnelModel,
	        FunnelForMaskedLM,
	        FunnelForPreTraining,
	        FunnelForQuestionAnswering,
	        FunnelForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	_UpperCamelCase					:					Any										=	(
	    {
	        'feature-extraction': (FunnelBaseModel, FunnelModel),
	        'fill-mask': FunnelForMaskedLM,
	        'question-answering': FunnelForQuestionAnswering,
	        'text-classification': FunnelForSequenceClassification,
	        'token-classification': FunnelForTokenClassification,
	        'zero-shot': FunnelForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case		,   snake_case		,   snake_case=False	):
		lowercase					  =   super()._prepare_for_class(snake_case		,   snake_case		,   return_labels=snake_case	)

		if return_labels:
			if model_class in get_values(snake_case	):
				lowercase					  =   torch.zeros(
				    (self.model_tester.batch_size, self.model_tester.seq_length)		,   dtype=torch.long		,   device=snake_case	)
		return inputs_dict





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   FunnelModelTester(self	)
		lowercase					  =   ConfigTester(self		,   config_class=snake_case	)





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





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_model(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_pretraining(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_masked_lm(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_token_classification(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_question_answering(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case	):




		if hasattr(snake_case		,   'weight'	) and module.weight is not None:
			module.weight.data.fill_(3	)
		if hasattr(snake_case		,   'bias'	) and module.bias is not None:
			module.bias.data.fill_(3	)

		for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
			if hasattr(snake_case		,   snake_case	) and getattr(snake_case		,   snake_case	) is not None:
				lowercase					  =   getattr(snake_case		,   snake_case	)
				weight.data.fill_(3	)







@require_torch
class 	A_       (      UpperCAmelCase__				,     unittest.TestCase       ):




	'''simple docstring'''


	_UpperCamelCase					:					List[str]										=	False
	_UpperCamelCase					:					str										=	False
	_UpperCamelCase					:					Tuple										=	(
	    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   FunnelModelTester(self		,   base=snake_case	)
		lowercase					  =   ConfigTester(self		,   config_class=snake_case	)





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





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_base_model(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_sequence_classification(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_multiple_choice(*snake_case	)





	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   self.model_tester.prepare_config_and_inputs_for_common()
		lowercase					  =   True

		for model_class in self.all_model_classes:
			if model_class.__name__ == "FunnelBaseModel":
				continue
			lowercase					  =   model_class(snake_case	)
			model.to(snake_case	)
			model.train()
			lowercase					  =   self._prepare_for_class(snake_case		,   snake_case		,   return_labels=snake_case	)
			lowercase					  =   model(**snake_case	).loss
			loss.backward()





	def      SCREAMING_SNAKE_CASE__						(	self		,   snake_case	):




		if hasattr(snake_case		,   'weight'	) and module.weight is not None:
			module.weight.data.fill_(3	)
		if hasattr(snake_case		,   'bias'	) and module.bias is not None:
			module.bias.data.fill_(3	)

		for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
			if hasattr(snake_case		,   snake_case	) and getattr(snake_case		,   snake_case	) is not None:
				lowercase					  =   getattr(snake_case		,   snake_case	)
				weight.data.fill_(3	)







@require_torch
@require_sentencepiece
@require_tokenizers
class 	A_       (      unittest.TestCase       ):




	'''simple docstring'''


	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   13
		lowercase					  =   7
		lowercase					  =   torch.arange(0		,   batch_size * sequence_length	).long().reshape(snake_case		,   snake_case	)
		lowercase					  =   [0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
		lowercase					  =   torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]	)

		lowercase					  =   FunnelModel.from_pretrained('sgugger/funnel-random-tiny'	)
		lowercase					  =   model(snake_case		,   token_type_ids=snake_case	)[0].abs()

		lowercase					  =   torch.tensor(2344.8352	)
		lowercase					  =   torch.tensor(0.8_052	)
		self.assertTrue(torch.allclose(output.sum()		,   snake_case		,   atol=1E-4	)	)
		self.assertTrue(torch.allclose(output.mean()		,   snake_case		,   atol=1E-4	)	)

		lowercase					  =   torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]	)
		lowercase					  =   model(snake_case		,   attention_mask=snake_case		,   token_type_ids=snake_case	)[0].abs()

		lowercase					  =   torch.tensor(2343.8425	)
		lowercase					  =   torch.tensor(0.8_049	)
		self.assertTrue(torch.allclose(output.sum()		,   snake_case		,   atol=1E-4	)	)
		self.assertTrue(torch.allclose(output.mean()		,   snake_case		,   atol=1E-4	)	)





	@slow
	def      SCREAMING_SNAKE_CASE__						(	self	):
		lowercase					  =   FunnelTokenizer.from_pretrained('huggingface/funnel-small'	)
		lowercase					  =   FunnelModel.from_pretrained('huggingface/funnel-small'	)
		lowercase					  =   tokenizer('Hello! I am the Funnel Transformer model.'		,   return_tensors='pt'	)
		lowercase					  =   model(**snake_case	)[0]

		lowercase					  =   torch.tensor(235.7246	)
		lowercase					  =   torch.tensor(0.0_256	)
		self.assertTrue(torch.allclose(output.sum()		,   snake_case		,   atol=1E-4	)	)
		self.assertTrue(torch.allclose(output.mean()		,   snake_case		,   atol=1E-4	)	)








import argparse
import json
from pathlib import Path

import requests
import torch
from datasets import load_dataset
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BeitConfig,
    BeitForImageClassification,
    BeitForMaskedImageModeling,
    BeitForSemanticSegmentation,
    BeitImageProcessor,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase     				=			logging.get_logger(__name__)






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE=False			,							__SCREAMING_SNAKE_CASE=False							):
	lowercase					  =   'backbone.' if is_semantic else ''

	lowercase					  =   []
	for i in range(config.num_hidden_layers							):
		# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
		rename_keys.append((F'''{prefix}blocks.{i}.norm1.weight''', F'''beit.encoder.layer.{i}.layernorm_before.weight''')							)
		rename_keys.append((F'''{prefix}blocks.{i}.norm1.bias''', F'''beit.encoder.layer.{i}.layernorm_before.bias''')							)
		rename_keys.append(
		    (F'''{prefix}blocks.{i}.attn.proj.weight''', F'''beit.encoder.layer.{i}.attention.output.dense.weight''')							)
		rename_keys.append(
		    (F'''{prefix}blocks.{i}.attn.proj.bias''', F'''beit.encoder.layer.{i}.attention.output.dense.bias''')							)
		rename_keys.append((F'''{prefix}blocks.{i}.norm2.weight''', F'''beit.encoder.layer.{i}.layernorm_after.weight''')							)
		rename_keys.append((F'''{prefix}blocks.{i}.norm2.bias''', F'''beit.encoder.layer.{i}.layernorm_after.bias''')							)
		rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc1.weight''', F'''beit.encoder.layer.{i}.intermediate.dense.weight''')							)
		rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc1.bias''', F'''beit.encoder.layer.{i}.intermediate.dense.bias''')							)
		rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc2.weight''', F'''beit.encoder.layer.{i}.output.dense.weight''')							)
		rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc2.bias''', F'''beit.encoder.layer.{i}.output.dense.bias''')							)

	# projection layer + position embeddings
	rename_keys.extend(
	    [
	        (F'''{prefix}cls_token''', 'beit.embeddings.cls_token'),
	        (F'''{prefix}patch_embed.proj.weight''', 'beit.embeddings.patch_embeddings.projection.weight'),
	        (F'''{prefix}patch_embed.proj.bias''', 'beit.embeddings.patch_embeddings.projection.bias'),
	    ]							)

	if has_lm_head:
		# mask token + shared relative position bias + layernorm
		rename_keys.extend(
		    [
		        ('mask_token', 'beit.embeddings.mask_token'),
		        (
		            'rel_pos_bias.relative_position_bias_table',
		            'beit.encoder.relative_position_bias.relative_position_bias_table',
		        ),
		        (
		            'rel_pos_bias.relative_position_index',
		            'beit.encoder.relative_position_bias.relative_position_index',
		        ),
		        ('norm.weight', 'layernorm.weight'),
		        ('norm.bias', 'layernorm.bias'),
		    ]							)
	elif is_semantic:
		# semantic segmentation classification heads
		rename_keys.extend(
		    [
		        ('decode_head.conv_seg.weight', 'decode_head.classifier.weight'),
		        ('decode_head.conv_seg.bias', 'decode_head.classifier.bias'),
		        ('auxiliary_head.conv_seg.weight', 'auxiliary_head.classifier.weight'),
		        ('auxiliary_head.conv_seg.bias', 'auxiliary_head.classifier.bias'),
		    ]							)
	else:
		# layernorm + classification head
		rename_keys.extend(
		    [
		        ('fc_norm.weight', 'beit.pooler.layernorm.weight'),
		        ('fc_norm.bias', 'beit.pooler.layernorm.bias'),
		        ('head.weight', 'classifier.weight'),
		        ('head.bias', 'classifier.bias'),
		    ]							)

	return rename_keys






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE=False			,							__SCREAMING_SNAKE_CASE=False							):
	for i in range(config.num_hidden_layers							):
		lowercase					  =   'backbone.' if is_semantic else ''
		# queries, keys and values
		lowercase					  =   state_dict.pop(F'''{prefix}blocks.{i}.attn.qkv.weight'''							)
		lowercase					  =   state_dict.pop(F'''{prefix}blocks.{i}.attn.q_bias'''							)
		lowercase					  =   state_dict.pop(F'''{prefix}blocks.{i}.attn.v_bias'''							)

		lowercase					  =   in_proj_weight[
		    : config.hidden_size, :
		]
		lowercase					  =   q_bias
		lowercase					  =   in_proj_weight[
		    config.hidden_size : config.hidden_size * 2, :
		]
		lowercase					  =   in_proj_weight[
		    -config.hidden_size :, :
		]
		lowercase					  =   v_bias

		# gamma_1 and gamma_2
		# we call them lambda because otherwise they are renamed when using .from_pretrained
		lowercase					  =   state_dict.pop(F'''{prefix}blocks.{i}.gamma_1'''							)
		lowercase					  =   state_dict.pop(F'''{prefix}blocks.{i}.gamma_2'''							)

		lowercase					  =   gamma_a
		lowercase					  =   gamma_a

		# relative_position bias table + index
		if not has_lm_head:
			# each layer has its own relative position bias
			lowercase					  =   state_dict.pop(F'''{prefix}blocks.{i}.attn.relative_position_bias_table'''							)
			lowercase					  =   state_dict.pop(F'''{prefix}blocks.{i}.attn.relative_position_index'''							)

			lowercase					  =   table
			lowercase					  =   index






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE							):
	lowercase					  =   dct.pop(__SCREAMING_SNAKE_CASE							)
	lowercase					  =   val






def 							UpperCAmelCase_		(							):
	lowercase					  =   'http://images.cocodataset.org/val2017/000000039769.jpg'
	lowercase					  =   Image.open(requests.get(__SCREAMING_SNAKE_CASE			,							stream=__SCREAMING_SNAKE_CASE							).raw							)
	return im






@torch.no_grad()
def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE							):
	lowercase					  =   BeitConfig()
	lowercase					  =   False
	lowercase					  =   False
	lowercase					  =   'huggingface/label-files'
	# set config parameters based on URL
	if checkpoint_url[-9:-4] == "pt22k":
		# masked image modeling
		lowercase					  =   True
		lowercase					  =   True
		lowercase					  =   True
	elif checkpoint_url[-9:-4] == "ft22k":
		# intermediate fine-tuning on ImageNet-22k
		lowercase					  =   True
		lowercase					  =   2_1841
		lowercase					  =   'imagenet-22k-id2label.json'
		lowercase					  =   json.load(open(hf_hub_download(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							repo_type='dataset'							)			,							'r'							)							)
		lowercase					  =   {int(__SCREAMING_SNAKE_CASE							): v for k, v in idalabel.items()}
		# this dataset contains 21843 labels but the model only has 21841
		# we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
		del idalabel[9205]
		del idalabel[1_5027]
		lowercase					  =   idalabel
		lowercase					  =   {v: k for k, v in idalabel.items()}
	elif checkpoint_url[-8:-4] == "to1k":
		# fine-tuning on ImageNet-1k
		lowercase					  =   True
		lowercase					  =   1000
		lowercase					  =   'imagenet-1k-id2label.json'
		lowercase					  =   json.load(open(hf_hub_download(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							repo_type='dataset'							)			,							'r'							)							)
		lowercase					  =   {int(__SCREAMING_SNAKE_CASE							): v for k, v in idalabel.items()}
		lowercase					  =   idalabel
		lowercase					  =   {v: k for k, v in idalabel.items()}
		if "384" in checkpoint_url:
			lowercase					  =   384
		if "512" in checkpoint_url:
			lowercase					  =   512
	elif "ade20k" in checkpoint_url:
		# fine-tuning
		lowercase					  =   True
		lowercase					  =   150
		lowercase					  =   'ade20k-id2label.json'
		lowercase					  =   json.load(open(hf_hub_download(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							repo_type='dataset'							)			,							'r'							)							)
		lowercase					  =   {int(__SCREAMING_SNAKE_CASE							): v for k, v in idalabel.items()}
		lowercase					  =   idalabel
		lowercase					  =   {v: k for k, v in idalabel.items()}
		lowercase					  =   640
		lowercase					  =   True
	else:
		raise ValueError('Checkpoint not supported, URL should either end with \'pt22k\', \'ft22k\', \'to1k\' or \'ade20k\''							)

	# size of the architecture
	if "base" in checkpoint_url:
		pass
	elif "large" in checkpoint_url:
		lowercase					  =   1024
		lowercase					  =   4096
		lowercase					  =   24
		lowercase					  =   16
		if "ade20k" in checkpoint_url:
			lowercase					  =   640
			lowercase					  =   [7, 11, 15, 23]
	else:
		raise ValueError('Should either find \'base\' or \'large\' in checkpoint URL'							)

	# load state_dict of original model, remove and rename some keys
	lowercase					  =   torch.hub.load_state_dict_from_url(__SCREAMING_SNAKE_CASE			,							map_location='cpu'			,							check_hash=__SCREAMING_SNAKE_CASE							)
	lowercase					  =   state_dict['model'] if 'ade20k' not in checkpoint_url else state_dict['state_dict']

	lowercase					  =   create_rename_keys(__SCREAMING_SNAKE_CASE			,							has_lm_head=__SCREAMING_SNAKE_CASE			,							is_semantic=__SCREAMING_SNAKE_CASE							)
	for src, dest in rename_keys:
		rename_key(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE							)
	read_in_q_k_v(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							has_lm_head=__SCREAMING_SNAKE_CASE			,							is_semantic=__SCREAMING_SNAKE_CASE							)
	if is_semantic:
		# add prefix to decoder keys
		for key, val in state_dict.copy().items():
			lowercase					  =   state_dict.pop(__SCREAMING_SNAKE_CASE							)
			if key.startswith('backbone.fpn'							):
				lowercase					  =   key.replace('backbone.fpn'			,							'fpn'							)
			lowercase					  =   val

    # load HuggingFace model
	if checkpoint_url[-9:-4] == "pt22k":
		lowercase					  =   BeitForMaskedImageModeling(__SCREAMING_SNAKE_CASE							)
	elif "ade20k" in checkpoint_url:
		lowercase					  =   BeitForSemanticSegmentation(__SCREAMING_SNAKE_CASE							)
	else:
		lowercase					  =   BeitForImageClassification(__SCREAMING_SNAKE_CASE							)
	model.eval()
	model.load_state_dict(__SCREAMING_SNAKE_CASE							)

	# Check outputs on an image
	if is_semantic:
		lowercase					  =   BeitImageProcessor(size=config.image_size			,							do_center_crop=__SCREAMING_SNAKE_CASE							)
		lowercase					  =   load_dataset('hf-internal-testing/fixtures_ade20k'			,							split='test'							)
		lowercase					  =   Image.open(ds[0]['file']							)
	else:
		lowercase					  =   BeitImageProcessor(
		    size=config.image_size			,							resample=PILImageResampling.BILINEAR			,							do_center_crop=__SCREAMING_SNAKE_CASE							)
		lowercase					  =   prepare_img()

	lowercase					  =   image_processor(images=__SCREAMING_SNAKE_CASE			,							return_tensors='pt'							)
	lowercase					  =   encoding['pixel_values']

	lowercase					  =   model(__SCREAMING_SNAKE_CASE							)
	lowercase					  =   outputs.logits

	# verify logits
	lowercase					  =   torch.Size([1, 1000]							)
	if checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k'							):
		lowercase					  =   torch.Size([1, 196, 8192]							)
	elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k'							):
		lowercase					  =   torch.Size([1, 196, 8192]							)
	elif checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k_ft22k'							):
		lowercase					  =   torch.Size([1, 2_1841]							)
		lowercase					  =   torch.tensor([2.22_88, 2.46_71, 0.73_95]							)
		lowercase					  =   2397
	elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k_ft22k'							):
		lowercase					  =   torch.Size([1, 2_1841]							)
		lowercase					  =   torch.tensor([1.68_81, -0.27_87, 0.59_01]							)
		lowercase					  =   2396
	elif checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k_ft1k'							):
		lowercase					  =   torch.tensor([0.12_41, 0.07_98, -0.65_69]							)
		lowercase					  =   285
	elif checkpoint_url[:-4].endswith('beit_base_patch16_224_pt22k_ft22kto1k'							):
		lowercase					  =   torch.tensor([-1.23_85, -1.09_87, -1.01_08]							)
		lowercase					  =   281
	elif checkpoint_url[:-4].endswith('beit_base_patch16_384_pt22k_ft22kto1k'							):
		lowercase					  =   torch.tensor([-1.53_03, -0.94_84, -0.31_47]							)
		lowercase					  =   761
	elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k_ft1k'							):
		lowercase					  =   torch.tensor([0.46_10, -0.09_28, 0.20_86]							)
		lowercase					  =   761
	elif checkpoint_url[:-4].endswith('beit_large_patch16_224_pt22k_ft22kto1k'							):
		lowercase					  =   torch.tensor([-0.48_04, 0.62_57, -0.18_37]							)
		lowercase					  =   761
	elif checkpoint_url[:-4].endswith('beit_large_patch16_384_pt22k_ft22kto1k'							):
		lowercase					  =   torch.tensor([[-0.51_22, 0.51_17, -0.21_13]]							)
		lowercase					  =   761
	elif checkpoint_url[:-4].endswith('beit_large_patch16_512_pt22k_ft22kto1k'							):
		lowercase					  =   torch.tensor([-0.30_62, 0.72_61, 0.48_52]							)
		lowercase					  =   761
	elif checkpoint_url[:-4].endswith('beit_base_patch16_640_pt22k_ft22ktoade20k'							):
		lowercase					  =   (1, 150, 160, 160)
		lowercase					  =   torch.tensor(
		    [
		        [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
		        [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
		        [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
		    ]							)
	elif checkpoint_url[:-4].endswith('beit_large_patch16_640_pt22k_ft22ktoade20k'							):
		lowercase					  =   (1, 150, 160, 160)
		lowercase					  =   torch.tensor(
		    [
		        [[-4.33_05, -2.30_49, -3.01_61], [-2.95_91, -1.53_05, -2.22_51], [-3.41_98, -1.80_04, -2.90_62]],
		        [[-5.89_22, -3.74_35, -4.39_78], [-4.20_63, -2.78_72, -3.47_55], [-4.27_91, -3.18_74, -4.16_81]],
		        [[0.98_95, 4.34_67, 4.76_63], [4.24_76, 5.68_30, 6.15_18], [4.55_50, 6.24_95, 6.51_54]],
		    ]							)
	else:
		raise ValueError('Can\'t verify logits as model is not supported'							)

	assert logits.shape == expected_shape, "Shape of logits not as expected"
	if not has_lm_head:
		if is_semantic:
			assert torch.allclose(
			    logits[0, :3, :3, :3]			,							__SCREAMING_SNAKE_CASE			,							atol=1e-3							), "First elements of logits not as expected"
		else:
			print('Predicted class idx:'			,							logits.argmax(-1							).item()							)
			assert torch.allclose(
			    logits[0, :3]			,							__SCREAMING_SNAKE_CASE			,							atol=1e-3							), "First elements of logits not as expected"
			assert logits.argmax(-1							).item() == expected_class_idx, "Predicted class index not as expected"

	Path(__SCREAMING_SNAKE_CASE							).mkdir(exist_ok=__SCREAMING_SNAKE_CASE							)
	print(F'''Saving model to {pytorch_dump_folder_path}'''							)
	model.save_pretrained(__SCREAMING_SNAKE_CASE							)
	print(F'''Saving image processor to {pytorch_dump_folder_path}'''							)
	image_processor.save_pretrained(__SCREAMING_SNAKE_CASE							)


if __name__ == "__main__":
	UpperCAmelCase     				=			argparse.ArgumentParser()

	parser.add_argument(
	    '''--checkpoint_url''',
	    default='''https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth''',
	    type=str,
	    help='''URL to the original PyTorch checkpoint (.pth file).''',
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
	)
	UpperCAmelCase     				=			parser.parse_args()
	convert_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)

