
'''simple docstring'''
import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, 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 EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
		from transformers.models.esm.modeling_esm import (
		    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    EsmEmbeddings,
		    create_position_ids_from_input_ids,
		)







class   a_ :




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase=13					,   UpperCAmelCase=7					,   UpperCAmelCase=False					,   UpperCAmelCase=True					,   UpperCAmelCase=False					,   UpperCAmelCase=True					,   UpperCAmelCase=33					,   UpperCAmelCase=32					,   UpperCAmelCase=5					,   UpperCAmelCase=4					,   UpperCAmelCase=37					,   UpperCAmelCase="gelu"					,   UpperCAmelCase=0.1					,   UpperCAmelCase=0.1					,   UpperCAmelCase=5_12					,   UpperCAmelCase=16					,   UpperCAmelCase=2					,   UpperCAmelCase=0.02					,   UpperCAmelCase=3					,   UpperCAmelCase=4					,   UpperCAmelCase=None					,   ):
				a_      				=						parent
				a_      				=						batch_size
				a_      				=						seq_length
				a_      				=						is_training
				a_      				=						use_input_mask
				a_      				=						use_token_type_ids
				a_      				=						use_labels
				a_      				=						vocab_size
				a_      				=						hidden_size
				a_      				=						num_hidden_layers
				a_      				=						num_attention_heads
				a_      				=						intermediate_size
				a_      				=						hidden_act
				a_      				=						hidden_dropout_prob
				a_      				=						attention_probs_dropout_prob
				a_      				=						max_position_embeddings
				a_      				=						type_vocab_size
				a_      				=						type_sequence_label_size
				a_      				=						initializer_range
				a_      				=						num_labels
				a_      				=						num_choices
				a_      				=						scope




			def 					lowerCAmelCase__					(       self	):
				a_      				=						ids_tensor([self.batch_size, self.seq_length]					,   self.vocab_size	)

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

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

				a_      				=						self.get_config()

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




			def 					lowerCAmelCase__					(       self	):
				return EsmConfig(
				    vocab_size=self.vocab_size					,   hidden_size=self.hidden_size					,   pad_token_id=1					,   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					,   )




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	):
				a_      				=						EsmModel(config=lowerCAmelCase__	)
				model.to(lowerCAmelCase__	)
				model.eval()
				a_      				=						model(lowerCAmelCase__					,   attention_mask=lowerCAmelCase__	)
				a_      				=						model(lowerCAmelCase__	)
				a_      				=						model(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					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	):
				a_      				=						EsmForMaskedLM(config=lowerCAmelCase__	)
				model.to(lowerCAmelCase__	)
				model.eval()
				a_      				=						model(lowerCAmelCase__					,   attention_mask=lowerCAmelCase__					,   labels=lowerCAmelCase__	)
				self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.vocab_size)	)




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	):
				a_      				=						self.num_labels
				a_      				=						EsmForTokenClassification(config=lowerCAmelCase__	)
				model.to(lowerCAmelCase__	)
				model.eval()
				a_      				=						model(lowerCAmelCase__					,   attention_mask=lowerCAmelCase__					,   labels=lowerCAmelCase__	)
				self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.num_labels)	)






			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.prepare_config_and_inputs()
				(
				    a_			
				)      				=						config_and_inputs
				a_      				=						{"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict







@require_torch
class   a_ (      a__		,  a__		,  unittest.TestCase					):
			lowerCamelCase__	:			Optional[Any]       =    False

			lowerCamelCase__	:			Tuple       =    (
			    (
			        EsmForMaskedLM,
			        EsmModel,
			        EsmForSequenceClassification,
			        EsmForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__	:			Tuple       =    ()
			lowerCamelCase__	:			Optional[int]       =    (
			    {
			        """feature-extraction""": EsmModel,
			        """fill-mask""": EsmForMaskedLM,
			        """text-classification""": EsmForSequenceClassification,
			        """token-classification""": EsmForTokenClassification,
			        """zero-shot""": EsmForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCamelCase__	:			Any       =    True




			def 					lowerCAmelCase__					(       self	):
				a_      				=						EsmModelTester(self	)
				a_      				=						ConfigTester(self					,   config_class=lowerCAmelCase__					,   hidden_size=37	)




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




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowerCAmelCase__	)




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




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__	)




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__	)




			@slow
			def 					lowerCAmelCase__					(       self	):
				for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					a_      				=						EsmModel.from_pretrained(lowerCAmelCase__	)
					self.assertIsNotNone(lowerCAmelCase__	)




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.model_tester.prepare_config_and_inputs()[0]
				a_      				=						EsmEmbeddings(config=lowerCAmelCase__	)

				a_      				=						torch.as_tensor([[12, 31, 13, model.padding_idx]]	)
				a_      				=						torch.as_tensor(
				    [
				        [
				            0 + model.padding_idx + 1,
				            1 + model.padding_idx + 1,
				            2 + model.padding_idx + 1,
				            model.padding_idx,
				        ]
				    ]	)
				a_      				=						create_position_ids_from_input_ids(lowerCAmelCase__					,   model.padding_idx	)
				self.assertEqual(position_ids.shape					,   expected_positions.shape	)
				self.assertTrue(torch.all(torch.eq(lowerCAmelCase__					,   lowerCAmelCase__	)	)	)




			def 					lowerCAmelCase__					(       self	):
				a_      				=						self.model_tester.prepare_config_and_inputs()[0]
				a_      				=						EsmEmbeddings(config=lowerCAmelCase__	)

				a_      				=						torch.empty(2					,   4					,   30	)
				a_      				=						[
				    0 + embeddings.padding_idx + 1,
				    1 + embeddings.padding_idx + 1,
				    2 + embeddings.padding_idx + 1,
				    3 + embeddings.padding_idx + 1,
				]
				a_      				=						torch.as_tensor([expected_single_positions, expected_single_positions]	)
				a_      				=						embeddings.create_position_ids_from_inputs_embeds(lowerCAmelCase__	)
				self.assertEqual(position_ids.shape					,   expected_positions.shape	)
				self.assertTrue(torch.all(torch.eq(lowerCAmelCase__					,   lowerCAmelCase__	)	)	)




			@unittest.skip("""Esm does not support embedding resizing"""	)
			def 					lowerCAmelCase__					(       self	):
				pass




			@unittest.skip("""Esm does not support embedding resizing"""	)
			def 					lowerCAmelCase__					(       self	):
				pass






			@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""	)
			def 					lowerCAmelCase__					(       self	):
				pass







@require_torch
class   a_ (      a__					):




			@slow
			def 					lowerCAmelCase__					(       self	):
				with torch.no_grad():
					a_      				=						EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""	)
					model.eval()
					a_      				=						torch.tensor([[0, 1, 2, 3, 4, 5]]	)
					a_      				=						model(lowerCAmelCase__	)[0]

					a_      				=						33

					a_      				=						torch.Size((1, 6, vocab_size)	)
					self.assertEqual(output.shape					,   lowerCAmelCase__	)

					a_      				=						torch.tensor(
					    [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]	)
					self.assertTrue(torch.allclose(output[:, :3, :3]					,   lowerCAmelCase__					,   atol=1e-4	)	)






			@slow
			def 					lowerCAmelCase__					(       self	):
				with torch.no_grad():
					a_      				=						EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""	)
					model.eval()

					a_      				=						torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]	)
					a_      				=						model(lowerCAmelCase__	)[0]
					# compare the actual values for a slice.
					a_      				=						torch.tensor(
					    [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]	)
					self.assertTrue(torch.allclose(output[:, :3, :3]					,   lowerCAmelCase__					,   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  _SCREAMING_SNAKE_CASE						(      _lowercase		:       List[Any]   )					->Dict:






	'''simple docstring'''






	a :						Any      =							[2, 2, 6, 2] if "tiny" in model_name else [2, 2, 18, 2]
	a :						str      =							True if "large" in model_name or "huge" in model_name else False
	a :						Optional[Any]      =							True if "large" in model_name or "huge" in model_name else False
	a :						Dict      =							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:
			a :						Union[str, Any]      =							[3, 3, 3, 3]
			a :						List[str]      =							[5, 5, 5, 5]
		elif "fl4" in model_name:
			a :						Any      =							[4, 4, 4, 4]
			a :						Optional[Any]      =							[3, 3, 3, 3]

	if "tiny" in model_name or "small" in model_name or "base" in model_name:
		a :						Dict      =							[3, 3, 3, 3]
		if "lrf" in model_name:
			a :						Optional[int]      =							[3, 3, 3, 3]
		else:
			a :						Tuple      =							[2, 2, 2, 2]

	if "tiny" in model_name:
		a :						List[str]      =							96
	elif "small" in model_name:
		a :						Union[str, Any]      =							96
	elif "base" in model_name:
		a :						Dict      =							128
	elif "large" in model_name:
		a :						Union[str, Any]      =							192
	elif "xlarge" in model_name:
		a :						Tuple      =							256
	elif "huge" in model_name:
		a :						List[str]      =							352

	# set label information
	a :						List[Any]      =							"huggingface/label-files"
	if "large" in model_name or "huge" in model_name:
		a :						Optional[int]      =							"imagenet-22k-id2label.json"
	else:
		a :						List[str]      =							"imagenet-1k-id2label.json"

	a :						Optional[int]      =							json.load(open(hf_hub_download(_lowercase       ,			_lowercase       ,			repo_type="dataset"   )       ,			"r"   )   )
	a :						str      =							{int(_lowercase   ): v for k, v in idalabel.items()}
	a :						List[str]      =							{v: k for k, v in idalabel.items()}

	a :						Dict      =							FocalNetConfig(
	    embed_dim=_lowercase       ,			depths=_lowercase       ,			focal_levels=_lowercase       ,			focal_windows=_lowercase       ,			use_conv_embed=_lowercase       ,			idalabel=_lowercase       ,			labelaid=_lowercase       ,			use_post_layernorm=_lowercase       ,			use_layerscale=_lowercase       ,			)

	return config






def  _SCREAMING_SNAKE_CASE						(      _lowercase		:       List[str]   )					->List[Any]:






	'''simple docstring'''






	if "patch_embed.proj" in name:
		a :						Any      =							name.replace("patch_embed.proj"       ,			"embeddings.patch_embeddings.projection"   )
	if "patch_embed.norm" in name:
		a :						List[str]      =							name.replace("patch_embed.norm"       ,			"embeddings.norm"   )
	if "layers" in name:
		a :						List[Any]      =							"encoder." + name
	if "encoder.layers" in name:
		a :						int      =							name.replace("encoder.layers"       ,			"encoder.stages"   )
	if "downsample.proj" in name:
		a :						Any      =							name.replace("downsample.proj"       ,			"downsample.projection"   )
	if "blocks" in name:
		a :						str      =							name.replace("blocks"       ,			"layers"   )
	if "modulation.f.weight" in name or "modulation.f.bias" in name:
		a :						Union[str, Any]      =							name.replace("modulation.f"       ,			"modulation.projection_in"   )
	if "modulation.h.weight" in name or "modulation.h.bias" in name:
		a :						Dict      =							name.replace("modulation.h"       ,			"modulation.projection_context"   )
	if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
		a :						Any      =							name.replace("modulation.proj"       ,			"modulation.projection_out"   )

	if name == "norm.weight":
		a :						str      =							"layernorm.weight"
	if name == "norm.bias":
		a :						Optional[Any]      =							"layernorm.bias"

	if "head" in name:
		a :						Tuple      =							name.replace("head"       ,			"classifier"   )
	else:
		a :						int      =							"focalnet." + name

	return name






def  _SCREAMING_SNAKE_CASE						(      _lowercase		:       Tuple       ,			_lowercase		:       Optional[Any]       ,			_lowercase		:       Tuple=False   )					->str:






	'''simple docstring'''






	a :						List[Any]      =							{
	    "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

	a :						str      =							model_name_to_url[model_name]
	print("Checkpoint URL: "       ,			_lowercase   )
	a :						Any      =							torch.hub.load_state_dict_from_url(_lowercase       ,			map_location="cpu"   )["model"]

	# rename keys
	for key in state_dict.copy().keys():
		a :						Any      =							state_dict.pop(_lowercase   )
		a :						Any      =							val

	a :						Any      =							get_focalnet_config(_lowercase   )
	a :						Optional[int]      =							FocalNetForImageClassification(_lowercase   )
	model.eval()

	# load state dict
	model.load_state_dict(_lowercase   )

	# verify conversion
	a :						Union[str, Any]      =							"http://images.cocodataset.org/val2017/000000039769.jpg"

	a :						Optional[int]      =							BitImageProcessor(
	    do_resize=_lowercase       ,			size={"shortest_edge": 256}       ,			resample=PILImageResampling.BILINEAR       ,			do_center_crop=_lowercase       ,			crop_size=224       ,			do_normalize=_lowercase       ,			image_mean=_lowercase       ,			image_std=_lowercase       ,			)
	a :						int      =							Image.open(requests.get(_lowercase       ,			stream=_lowercase   ).raw   )
	a :						Dict      =							processor(images=_lowercase       ,			return_tensors="pt"   )

	a :						Optional[int]      =							transforms.Compose(
	    [
	        transforms.Resize(256   ),
	        transforms.CenterCrop(224   ),
	        transforms.ToTensor(),
	        transforms.Normalize(mean=[0.485, 0.456, 0.406]       ,			std=[0.229, 0.224, 0.225]   ),
	    ]   )

	a :						str      =							image_transforms(_lowercase   ).unsqueeze(0   )

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

	a :						Dict      =							model(**_lowercase   )

	a :						List[str]      =							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":
		a :						Union[str, Any]      =							torch.tensor([0.2166, -0.4368, 0.2191]   )
	elif model_name == "focalnet-tiny-lrf":
		a :						Union[str, Any]      =							torch.tensor([1.1669, 0.0125, -0.1695]   )
	elif model_name == "focalnet-small":
		a :						Dict      =							torch.tensor([0.4917, -0.0430, 0.1341]   )
	elif model_name == "focalnet-small-lrf":
		a :						Dict      =							torch.tensor([-0.2588, -0.5342, -0.2331]   )
	elif model_name == "focalnet-base":
		a :						Any      =							torch.tensor([-0.1655, -0.4090, -0.1730]   )
	elif model_name == "focalnet-base-lrf":
		a :						str      =							torch.tensor([0.5306, -0.0483, -0.3928]   )
	assert torch.allclose(outputs.logits[0, :3]       ,			_lowercase       ,			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(_lowercase   )
		processor.save_pretrained(_lowercase   )

	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__":
		a		:					str	      =     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.''',
		)

		a		:					Tuple	      =     parser.parse_args()
		convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)


