
"""simple docstring"""




import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
						import torch

						from transformers import (
						    MODEL_FOR_PRETRAINING_MAPPING,
						    AlbertForMaskedLM,
						    AlbertForMultipleChoice,
						    AlbertForPreTraining,
						    AlbertForQuestionAnswering,
						    AlbertForSequenceClassification,
						    AlbertForTokenClassification,
						    AlbertModel,
						)
						from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST





class  snake_case     :

							def __init__(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_=13       ,       SCREAMING_SNAKE_CASE_=7       ,       SCREAMING_SNAKE_CASE_=True       ,       SCREAMING_SNAKE_CASE_=True       ,       SCREAMING_SNAKE_CASE_=True       ,       SCREAMING_SNAKE_CASE_=True       ,       SCREAMING_SNAKE_CASE_=99       ,       SCREAMING_SNAKE_CASE_=16       ,       SCREAMING_SNAKE_CASE_=36       ,       SCREAMING_SNAKE_CASE_=6       ,       SCREAMING_SNAKE_CASE_=6       ,       SCREAMING_SNAKE_CASE_=6       ,       SCREAMING_SNAKE_CASE_=37       ,       SCREAMING_SNAKE_CASE_="gelu"       ,       SCREAMING_SNAKE_CASE_=0.1       ,       SCREAMING_SNAKE_CASE_=0.1       ,       SCREAMING_SNAKE_CASE_=5_12       ,       SCREAMING_SNAKE_CASE_=16       ,       SCREAMING_SNAKE_CASE_=2       ,       SCREAMING_SNAKE_CASE_=0.02       ,       SCREAMING_SNAKE_CASE_=3       ,       SCREAMING_SNAKE_CASE_=4       ,       SCREAMING_SNAKE_CASE_=None       ,       ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = parent
									SCREAMING_SNAKE_CASE_						 = batch_size
									SCREAMING_SNAKE_CASE_						 = seq_length
									SCREAMING_SNAKE_CASE_						 = is_training
									SCREAMING_SNAKE_CASE_						 = use_input_mask
									SCREAMING_SNAKE_CASE_						 = use_token_type_ids
									SCREAMING_SNAKE_CASE_						 = use_labels
									SCREAMING_SNAKE_CASE_						 = vocab_size
									SCREAMING_SNAKE_CASE_						 = embedding_size
									SCREAMING_SNAKE_CASE_						 = hidden_size
									SCREAMING_SNAKE_CASE_						 = num_hidden_layers
									SCREAMING_SNAKE_CASE_						 = num_hidden_groups
									SCREAMING_SNAKE_CASE_						 = num_attention_heads
									SCREAMING_SNAKE_CASE_						 = intermediate_size
									SCREAMING_SNAKE_CASE_						 = hidden_act
									SCREAMING_SNAKE_CASE_						 = hidden_dropout_prob
									SCREAMING_SNAKE_CASE_						 = attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE_						 = max_position_embeddings
									SCREAMING_SNAKE_CASE_						 = type_vocab_size
									SCREAMING_SNAKE_CASE_						 = type_sequence_label_size
									SCREAMING_SNAKE_CASE_						 = initializer_range
									SCREAMING_SNAKE_CASE_						 = num_labels
									SCREAMING_SNAKE_CASE_						 = num_choices
									SCREAMING_SNAKE_CASE_						 = scope

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = ids_tensor([self.batch_size, self.seq_length]       ,       self.vocab_size    )

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

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

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

									SCREAMING_SNAKE_CASE_						 = self.get_config()

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

							def _lowercase					(self    ):

									"""simple docstring"""







									return AlbertConfig(
									    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       ,       num_hidden_groups=self.num_hidden_groups       ,       )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = AlbertModel(config=lowercase_    )
									model.to(lowercase_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(lowercase_       ,       attention_mask=lowercase_       ,       token_type_ids=lowercase_    )
									SCREAMING_SNAKE_CASE_						 = model(lowercase_       ,       token_type_ids=lowercase_    )
									SCREAMING_SNAKE_CASE_						 = model(lowercase_    )
									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 _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = AlbertForPreTraining(config=lowercase_    )
									model.to(lowercase_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(
									    lowercase_       ,       attention_mask=lowercase_       ,       token_type_ids=lowercase_       ,       labels=lowercase_       ,       sentence_order_label=lowercase_       ,       )
									self.parent.assertEqual(result.prediction_logits.shape       ,       (self.batch_size, self.seq_length, self.vocab_size)    )
									self.parent.assertEqual(result.sop_logits.shape       ,       (self.batch_size, config.num_labels)    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = AlbertForMaskedLM(config=lowercase_    )
									model.to(lowercase_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(lowercase_       ,       attention_mask=lowercase_       ,       token_type_ids=lowercase_       ,       labels=lowercase_    )
									self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.seq_length, self.vocab_size)    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = AlbertForQuestionAnswering(config=lowercase_    )
									model.to(lowercase_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(
									    lowercase_       ,       attention_mask=lowercase_       ,       token_type_ids=lowercase_       ,       start_positions=lowercase_       ,       end_positions=lowercase_       ,       )
									self.parent.assertEqual(result.start_logits.shape       ,       (self.batch_size, self.seq_length)    )
									self.parent.assertEqual(result.end_logits.shape       ,       (self.batch_size, self.seq_length)    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.num_labels
									SCREAMING_SNAKE_CASE_						 = AlbertForSequenceClassification(lowercase_    )
									model.to(lowercase_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(lowercase_       ,       attention_mask=lowercase_       ,       token_type_ids=lowercase_       ,       labels=lowercase_    )
									self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.num_labels)    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.num_labels
									SCREAMING_SNAKE_CASE_						 = AlbertForTokenClassification(config=lowercase_    )
									model.to(lowercase_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(lowercase_       ,       attention_mask=lowercase_       ,       token_type_ids=lowercase_       ,       labels=lowercase_    )
									self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.seq_length, self.num_labels)    )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.num_choices
									SCREAMING_SNAKE_CASE_						 = AlbertForMultipleChoice(config=lowercase_    )
									model.to(lowercase_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = input_ids.unsqueeze(1    ).expand(-1       ,       self.num_choices       ,       -1    ).contiguous()
									SCREAMING_SNAKE_CASE_						 = token_type_ids.unsqueeze(1    ).expand(-1       ,       self.num_choices       ,       -1    ).contiguous()
									SCREAMING_SNAKE_CASE_						 = input_mask.unsqueeze(1    ).expand(-1       ,       self.num_choices       ,       -1    ).contiguous()
									SCREAMING_SNAKE_CASE_						 = model(
									    lowercase_       ,       attention_mask=lowercase_       ,       token_type_ids=lowercase_       ,       labels=lowercase_       ,       )
									self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.num_choices)    )







							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.prepare_config_and_inputs()
									(
									    SCREAMING_SNAKE_CASE_ 
									)						 = config_and_inputs
									SCREAMING_SNAKE_CASE_						 = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
									return config, inputs_dict





@require_torch
class  snake_case     (							snake_case__       ,			snake_case__       ,			unittest.TestCase   ):
							UpperCAmelCase__						     =      (
							    (
							        AlbertModel,
							        AlbertForPreTraining,
							        AlbertForMaskedLM,
							        AlbertForMultipleChoice,
							        AlbertForSequenceClassification,
							        AlbertForTokenClassification,
							        AlbertForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							UpperCAmelCase__						     =      (
							    {
							        '''feature-extraction''': AlbertModel,
							        '''fill-mask''': AlbertForMaskedLM,
							        '''question-answering''': AlbertForQuestionAnswering,
							        '''text-classification''': AlbertForSequenceClassification,
							        '''token-classification''': AlbertForTokenClassification,
							        '''zero-shot''': AlbertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							UpperCAmelCase__						     =      True

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_=False    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = super()._prepare_for_class(lowercase_       ,       lowercase_       ,       return_labels=lowercase_    )

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

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = AlbertModelTester(self    )
									SCREAMING_SNAKE_CASE_						 = ConfigTester(self       ,       config_class=lowercase_       ,       hidden_size=37    )

							def _lowercase					(self    ):

									"""simple docstring"""







									self.config_tester.run_common_tests()

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowercase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*lowercase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*lowercase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*lowercase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*lowercase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*lowercase_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
											SCREAMING_SNAKE_CASE_						 = type
											self.model_tester.create_and_check_model(*lowercase_    )







							@slow
							def _lowercase					(self    ):

									"""simple docstring"""







									for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											SCREAMING_SNAKE_CASE_						 = AlbertModel.from_pretrained(lowercase_    )
											self.assertIsNotNone(lowercase_    )





@require_torch
class  snake_case     (							unittest.TestCase   ):

							@slow
							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = AlbertModel.from_pretrained('''albert-base-v2'''    )
									SCREAMING_SNAKE_CASE_						 = torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]    )
									SCREAMING_SNAKE_CASE_						 = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
									with torch.no_grad():
											SCREAMING_SNAKE_CASE_						 = model(lowercase_       ,       attention_mask=lowercase_    )[0]
									SCREAMING_SNAKE_CASE_						 = torch.Size((1, 11, 7_68)    )
									self.assertEqual(output.shape       ,       lowercase_    )
									SCREAMING_SNAKE_CASE_						 = torch.tensor(
									    [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]    )

									self.assertTrue(torch.allclose(output[:, 1:4, 1:4]       ,       lowercase_       ,       atol=1e-4    )    )

"""simple docstring"""




import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling



def        _lowerCamelCase ( __a    ):
		SCREAMING_SNAKE_CASE_						 = [2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
		SCREAMING_SNAKE_CASE_						 = True if '''large''' in model_name or '''huge''' in model_name else False
		SCREAMING_SNAKE_CASE_						 = True if '''large''' in model_name or '''huge''' in model_name else False
		SCREAMING_SNAKE_CASE_						 = True if '''large''' in model_name or '''huge''' in model_name else False

		if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
				if "fl3" in model_name:
						SCREAMING_SNAKE_CASE_						 = [3, 3, 3, 3]
						SCREAMING_SNAKE_CASE_						 = [5, 5, 5, 5]
				elif "fl4" in model_name:
						SCREAMING_SNAKE_CASE_						 = [4, 4, 4, 4]
						SCREAMING_SNAKE_CASE_						 = [3, 3, 3, 3]

		if "tiny" in model_name or "small" in model_name or "base" in model_name:
				SCREAMING_SNAKE_CASE_						 = [3, 3, 3, 3]
				if "lrf" in model_name:
						SCREAMING_SNAKE_CASE_						 = [3, 3, 3, 3]
				else:
						SCREAMING_SNAKE_CASE_						 = [2, 2, 2, 2]

		if "tiny" in model_name:
				SCREAMING_SNAKE_CASE_						 = 96
		elif "small" in model_name:
				SCREAMING_SNAKE_CASE_						 = 96
		elif "base" in model_name:
				SCREAMING_SNAKE_CASE_						 = 128
		elif "large" in model_name:
				SCREAMING_SNAKE_CASE_						 = 192
		elif "xlarge" in model_name:
				SCREAMING_SNAKE_CASE_						 = 256
		elif "huge" in model_name:
				SCREAMING_SNAKE_CASE_						 = 352

		# set label information
		SCREAMING_SNAKE_CASE_						 = '''huggingface/label-files'''
		if "large" in model_name or "huge" in model_name:
				SCREAMING_SNAKE_CASE_						 = '''imagenet-22k-id2label.json'''
		else:
				SCREAMING_SNAKE_CASE_						 = '''imagenet-1k-id2label.json'''

		SCREAMING_SNAKE_CASE_						 = json.load(open(hf_hub_download(__a,							__a,							repo_type='''dataset'''    ),							'''r'''    )    )
		SCREAMING_SNAKE_CASE_						 = {int(__a    ): v for k, v in idalabel.items()}
		SCREAMING_SNAKE_CASE_						 = {v: k for k, v in idalabel.items()}

		SCREAMING_SNAKE_CASE_						 = FocalNetConfig(
		    embed_dim=__a,							depths=__a,							focal_levels=__a,							focal_windows=__a,							use_conv_embed=__a,							idalabel=__a,							labelaid=__a,							use_post_layernorm=__a,							use_layerscale=__a,							)

		return config



def        _lowerCamelCase ( __a    ):
		if "patch_embed.proj" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''patch_embed.proj''',							'''embeddings.patch_embeddings.projection'''    )
		if "patch_embed.norm" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''patch_embed.norm''',							'''embeddings.norm'''    )
		if "layers" in name:
				SCREAMING_SNAKE_CASE_						 = '''encoder.''' + name
		if "encoder.layers" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''encoder.layers''',							'''encoder.stages'''    )
		if "downsample.proj" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''downsample.proj''',							'''downsample.projection'''    )
		if "blocks" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''blocks''',							'''layers'''    )
		if "modulation.f.weight" in name or "modulation.f.bias" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''modulation.f''',							'''modulation.projection_in'''    )
		if "modulation.h.weight" in name or "modulation.h.bias" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''modulation.h''',							'''modulation.projection_context'''    )
		if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''modulation.proj''',							'''modulation.projection_out'''    )

		if name == "norm.weight":
				SCREAMING_SNAKE_CASE_						 = '''layernorm.weight'''
		if name == "norm.bias":
				SCREAMING_SNAKE_CASE_						 = '''layernorm.bias'''

		if "head" in name:
				SCREAMING_SNAKE_CASE_						 = name.replace('''head''',							'''classifier'''    )
		else:
				SCREAMING_SNAKE_CASE_						 = '''focalnet.''' + name

		return name



def        _lowerCamelCase ( __a,							__a,							__a=False    ):
		# fmt: off
		SCREAMING_SNAKE_CASE_						 = {
		    '''focalnet-tiny''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth''',
		    '''focalnet-tiny-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth''',
		    '''focalnet-small''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth''',
		    '''focalnet-small-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth''',
		    '''focalnet-base''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth''',
		    '''focalnet-base-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth''',
		    '''focalnet-large-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth''',
		    '''focalnet-large-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth''',
		    '''focalnet-xlarge-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth''',
		    '''focalnet-xlarge-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth''',
		}
		# fmt: on

		SCREAMING_SNAKE_CASE_						 = model_name_to_url[model_name]
		print('''Checkpoint URL: ''',							__a    )
		SCREAMING_SNAKE_CASE_						 = torch.hub.load_state_dict_from_url(__a,							map_location='''cpu'''    )['''model''']

		# rename keys
		for key in state_dict.copy().keys():
				SCREAMING_SNAKE_CASE_						 = state_dict.pop(__a    )
				SCREAMING_SNAKE_CASE_						 = val

		SCREAMING_SNAKE_CASE_						 = get_focalnet_config(__a    )
		SCREAMING_SNAKE_CASE_						 = FocalNetForImageClassification(__a    )
		model.eval()

		# load state dict
		model.load_state_dict(__a    )

		# verify conversion
		SCREAMING_SNAKE_CASE_						 = '''http://images.cocodataset.org/val2017/000000039769.jpg'''

		SCREAMING_SNAKE_CASE_						 = BitImageProcessor(
		    do_resize=__a,							size={'''shortest_edge''': 256},							resample=PILImageResampling.BILINEAR,							do_center_crop=__a,							crop_size=224,							do_normalize=__a,							image_mean=__a,							image_std=__a,							)
		SCREAMING_SNAKE_CASE_						 = Image.open(requests.get(__a,							stream=__a    ).raw    )
		SCREAMING_SNAKE_CASE_						 = processor(images=__a,							return_tensors='''pt'''    )

		SCREAMING_SNAKE_CASE_						 = transforms.Compose(
		    [
		        transforms.Resize(256    ),
		        transforms.CenterCrop(224    ),
		        transforms.ToTensor(),
		        transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6],							std=[0.2_2_9, 0.2_2_4, 0.2_2_5]    ),
		    ]    )

		SCREAMING_SNAKE_CASE_						 = image_transforms(__a    ).unsqueeze(0    )

		# verify pixel_values
		assert torch.allclose(inputs.pixel_values,							__a,							atol=1E-4    )

		SCREAMING_SNAKE_CASE_						 = model(**__a    )

		SCREAMING_SNAKE_CASE_						 = outputs.logits.argmax(-1    ).item()
		print('''Predicted class:''',							model.config.idalabel[predicted_class_idx]    )

		print('''First values of logits:''',							outputs.logits[0, :3]    )

		if model_name == "focalnet-tiny":
				SCREAMING_SNAKE_CASE_						 = torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]    )
		elif model_name == "focalnet-tiny-lrf":
				SCREAMING_SNAKE_CASE_						 = torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]    )
		elif model_name == "focalnet-small":
				SCREAMING_SNAKE_CASE_						 = torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]    )
		elif model_name == "focalnet-small-lrf":
				SCREAMING_SNAKE_CASE_						 = torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]    )
		elif model_name == "focalnet-base":
				SCREAMING_SNAKE_CASE_						 = torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]    )
		elif model_name == "focalnet-base-lrf":
				SCREAMING_SNAKE_CASE_						 = torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]    )
		assert torch.allclose(outputs.logits[0, :3],							__a,							atol=1E-4    )
		print('''Looks ok!'''    )

		if pytorch_dump_folder_path is not None:
				print(F'Saving model and processor of {model_name} to {pytorch_dump_folder_path}'    )
				model.save_pretrained(__a    )
				processor.save_pretrained(__a    )

		if push_to_hub:
				print(F'Pushing model and processor of {model_name} to the hub...'    )
				model.push_to_hub(F'{model_name}'    )
				processor.push_to_hub(F'{model_name}'    )


if __name__ == "__main__":
						lowerCAmelCase__							=				argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--model_name',
						    default='focalnet-tiny',
						    type=str,
						    help='Name of the FocalNet model you\'d like to convert.',
						)
						parser.add_argument(
						    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
						)
						parser.add_argument(
						    '--push_to_hub',
						    action='store_true',
						    help='Whether to push the model and processor to the hub.',
						)

						lowerCAmelCase__							=				parser.parse_args()
						convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)