'''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 						__A							(     UpperCAmelCase					)     ->			Tuple:




						'''simple docstring'''
						_UpperCamelCase			:			Dict   			=						[2, 2, 6, 2] if "tiny" in model_name else [2, 2, 1_8, 2]
						_UpperCamelCase			:			Any   			=						True if "large" in model_name or "huge" in model_name else False
						_UpperCamelCase			:			Any   			=						True if "large" in model_name or "huge" in model_name else False
						_UpperCamelCase			:			List[Any]   			=						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:
																		_UpperCamelCase			:			Optional[Any]   			=						[3, 3, 3, 3]
																		_UpperCamelCase			:			Tuple   			=						[5, 5, 5, 5]
												elif "fl4" in model_name:
																		_UpperCamelCase			:			Dict   			=						[4, 4, 4, 4]
																		_UpperCamelCase			:			List[str]   			=						[3, 3, 3, 3]

						if "tiny" in model_name or "small" in model_name or "base" in model_name:
												_UpperCamelCase			:			Union[str, Any]   			=						[3, 3, 3, 3]
												if "lrf" in model_name:
																		_UpperCamelCase			:			Optional[int]   			=						[3, 3, 3, 3]
												else:
																		_UpperCamelCase			:			Any   			=						[2, 2, 2, 2]

						if "tiny" in model_name:
												_UpperCamelCase			:			Union[str, Any]   			=						9_6
						elif "small" in model_name:
												_UpperCamelCase			:			Tuple   			=						9_6
						elif "base" in model_name:
												_UpperCamelCase			:			Any   			=						1_2_8
						elif "large" in model_name:
												_UpperCamelCase			:			str   			=						1_9_2
						elif "xlarge" in model_name:
												_UpperCamelCase			:			Any   			=						2_5_6
						elif "huge" in model_name:
												_UpperCamelCase			:			Dict   			=						3_5_2

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

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

						_UpperCamelCase			:			Any   			=						FocalNetConfig(
						    embed_dim=UpperCAmelCase     ,depths=UpperCAmelCase     ,focal_levels=UpperCAmelCase     ,focal_windows=UpperCAmelCase     ,use_conv_embed=UpperCAmelCase     ,idalabel=UpperCAmelCase     ,labelaid=UpperCAmelCase     ,use_post_layernorm=UpperCAmelCase     ,use_layerscale=UpperCAmelCase     ,)

						return config
def 						__A							(     UpperCAmelCase					)     ->			List[str]:




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

						if name == "norm.weight":
												_UpperCamelCase			:			Any   			=						"layernorm.weight"
						if name == "norm.bias":
												_UpperCamelCase			:			int   			=						"layernorm.bias"

						if "head" in name:
												_UpperCamelCase			:			Dict   			=						name.replace("head"     ,"classifier"					)
						else:
												_UpperCamelCase			:			List[Any]   			=						"focalnet." + name

						return name
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase=False					)     ->			Optional[Any]:




						'''simple docstring'''
						# fmt: off
						_UpperCamelCase			:			List[str]   			=						{
						    "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

						_UpperCamelCase			:			Tuple   			=						model_name_to_url[model_name]
						print("Checkpoint URL: "     ,UpperCAmelCase					)
						_UpperCamelCase			:			Union[str, Any]   			=						torch.hub.load_state_dict_from_url(UpperCAmelCase     ,map_location="cpu"					)["model"]

						# rename keys
						for key in state_dict.copy().keys():
												_UpperCamelCase			:			Optional[Any]   			=						state_dict.pop(UpperCAmelCase					)
												_UpperCamelCase			:			Optional[int]   			=						val

						_UpperCamelCase			:			Any   			=						get_focalnet_config(UpperCAmelCase					)
						_UpperCamelCase			:			int   			=						FocalNetForImageClassification(UpperCAmelCase					)
						model.eval()

						# load state dict
						model.load_state_dict(UpperCAmelCase					)

						# verify conversion
						_UpperCamelCase			:			Optional[Any]   			=						"http://images.cocodataset.org/val2017/000000039769.jpg"

						_UpperCamelCase			:			Optional[int]   			=						BitImageProcessor(
						    do_resize=UpperCAmelCase     ,size={"shortest_edge": 2_5_6}     ,resample=PILImageResampling.BILINEAR     ,do_center_crop=UpperCAmelCase     ,crop_size=2_2_4     ,do_normalize=UpperCAmelCase     ,image_mean=UpperCAmelCase     ,image_std=UpperCAmelCase     ,)
						_UpperCamelCase			:			List[str]   			=						Image.open(requests.get(UpperCAmelCase     ,stream=UpperCAmelCase					).raw					)
						_UpperCamelCase			:			int   			=						processor(images=UpperCAmelCase     ,return_tensors="pt"					)

						_UpperCamelCase			:			List[str]   			=						transforms.Compose(
						    [
						        transforms.Resize(2_5_6					),
						        transforms.CenterCrop(2_2_4					),
						        transforms.ToTensor(),
						        transforms.Normalize(mean=[0.4_85, 0.4_56, 0.4_06]     ,std=[0.2_29, 0.2_24, 0.2_25]					),
						    ]					)

						_UpperCamelCase			:			List[str]   			=						image_transforms(UpperCAmelCase					).unsqueeze(0					)

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

						_UpperCamelCase			:			Any   			=						model(**UpperCAmelCase					)

						_UpperCamelCase			:			Any   			=						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":
												_UpperCamelCase			:			Tuple   			=						torch.tensor([0.21_66, -0.43_68, 0.21_91]					)
						elif model_name == "focalnet-tiny-lrf":
												_UpperCamelCase			:			List[Any]   			=						torch.tensor([1.16_69, 0.01_25, -0.16_95]					)
						elif model_name == "focalnet-small":
												_UpperCamelCase			:			Tuple   			=						torch.tensor([0.49_17, -0.04_30, 0.13_41]					)
						elif model_name == "focalnet-small-lrf":
												_UpperCamelCase			:			int   			=						torch.tensor([-0.25_88, -0.53_42, -0.23_31]					)
						elif model_name == "focalnet-base":
												_UpperCamelCase			:			List[str]   			=						torch.tensor([-0.16_55, -0.40_90, -0.17_30]					)
						elif model_name == "focalnet-base-lrf":
												_UpperCamelCase			:			Optional[Any]   			=						torch.tensor([0.53_06, -0.04_83, -0.39_28]					)
						assert torch.allclose(outputs.logits[0, :3]     ,UpperCAmelCase     ,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(UpperCAmelCase					)
												processor.save_pretrained(UpperCAmelCase					)

						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_	:		Tuple			=     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_	:		str			=     parser.parse_args()
			convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



'''simple docstring'''





from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


lowerCAmelCase_	:		int			=     logging.get_logger(__name__)

lowerCAmelCase_	:		Any			=     {
    """EleutherAI/gpt-j-6B""": """https://huggingface.co/EleutherAI/gpt-j-6B/resolve/main/config.json""",
    # See all GPT-J models at https://huggingface.co/models?filter=gpt_j
}




class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''
				UpperCAmelCase__							=							'''gptj'''
				UpperCAmelCase__							=							{
				    '''max_position_embeddings''': '''n_positions''',
				    '''hidden_size''': '''n_embd''',
				    '''num_attention_heads''': '''n_head''',
				    '''num_hidden_layers''': '''n_layer''',
				}





				def __init__(						self					:						Optional[Any]       ,		lowercase__					:						Union[str, Any]=50_400       ,		lowercase__					:						Union[str, Any]=2_048       ,		lowercase__					:						Tuple=4_096       ,		lowercase__					:						List[str]=28       ,		lowercase__					:						Optional[int]=16       ,		lowercase__					:						str=64       ,		lowercase__					:						Any=None       ,		lowercase__					:						Any="gelu_new"       ,		lowercase__					:						Union[str, Any]=0.0       ,		lowercase__					:						Optional[Any]=0.0       ,		lowercase__					:						Any=0.0       ,		lowercase__					:						Tuple=1e-5       ,		lowercase__					:						Any=0.0_2       ,		lowercase__					:						int=True       ,		lowercase__					:						int=50_256       ,		lowercase__					:						Any=50_256       ,		lowercase__					:						Tuple=False       ,		**lowercase__					:						str       ,		)							->Optional[Any]:


										'''simple docstring'''




										_UpperCamelCase			:			Dict   			=						vocab_size
										_UpperCamelCase			:			List[str]   			=						n_positions
										_UpperCamelCase			:			Union[str, Any]   			=						n_embd
										_UpperCamelCase			:			Union[str, Any]   			=						n_layer
										_UpperCamelCase			:			Optional[Any]   			=						n_head
										_UpperCamelCase			:			Dict   			=						n_inner
										_UpperCamelCase			:			Optional[Any]   			=						rotary_dim
										_UpperCamelCase			:			Tuple   			=						activation_function
										_UpperCamelCase			:			List[Any]   			=						resid_pdrop
										_UpperCamelCase			:			Any   			=						embd_pdrop
										_UpperCamelCase			:			Optional[Any]   			=						attn_pdrop
										_UpperCamelCase			:			Optional[Any]   			=						layer_norm_epsilon
										_UpperCamelCase			:			Union[str, Any]   			=						initializer_range
										_UpperCamelCase			:			Optional[int]   			=						use_cache

										_UpperCamelCase			:			str   			=						bos_token_id
										_UpperCamelCase			:			Any   			=						eos_token_id

										super().__init__(
										    bos_token_id=lowercase__       ,		eos_token_id=lowercase__       ,		tie_word_embeddings=lowercase__       ,		**lowercase__					)





class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''





				def __init__(						self					:						Optional[int]       ,		lowercase__					:						PretrainedConfig       ,		lowercase__					:						str = "default"       ,		lowercase__					:						List[PatchingSpec] = None       ,		lowercase__					:						bool = False       ,		)							->Union[str, Any]:


										'''simple docstring'''




										super().__init__(lowercase__       ,		task=lowercase__       ,		patching_specs=lowercase__       ,		use_past=lowercase__					)
										if not getattr(self._config       ,		"pad_token_id"       ,		lowercase__					):
																# TODO: how to do that better?
																_UpperCamelCase			:			Optional[int]   			=						0





				@property
				def  snake_case__				(						self					:						List[str]					)							->Mapping[str, Mapping[int, str]]:


										'''simple docstring'''




										_UpperCamelCase			:			List[str]   			=						OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}					)
										if self.use_past:
																self.fill_with_past_key_values_(lowercase__       ,		direction="inputs"					)
																_UpperCamelCase			:			str   			=						{0: "batch", 1: "past_sequence + sequence"}
										else:
																_UpperCamelCase			:			Optional[Any]   			=						{0: "batch", 1: "sequence"}

										return common_inputs





				@property
				def  snake_case__				(						self					:						int					)							->int:


										'''simple docstring'''




										return self._config.n_layer





				@property
				def  snake_case__				(						self					:						Dict					)							->int:


										'''simple docstring'''




										return self._config.n_head





				def  snake_case__				(						self					:						int       ,		lowercase__					:						PreTrainedTokenizer       ,		lowercase__					:						int = -1       ,		lowercase__					:						int = -1       ,		lowercase__					:						bool = False       ,		lowercase__					:						Optional[TensorType] = None       ,		)							->Mapping[str, Any]:


										'''simple docstring'''




										_UpperCamelCase			:			int   			=						super(lowercase__       ,		self					).generate_dummy_inputs(
										    lowercase__       ,		batch_size=lowercase__       ,		seq_length=lowercase__       ,		is_pair=lowercase__       ,		framework=lowercase__					)

										# We need to order the input in the way they appears in the forward()
										_UpperCamelCase			:			Tuple   			=						OrderedDict({"input_ids": common_inputs["input_ids"]}					)

										# Need to add the past_keys
										if self.use_past:
																if not is_torch_available():
																						raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."					)
																else:
																						import torch

																						_UpperCamelCase			,       _UpperCamelCase			:			Optional[Any]   			=						common_inputs["input_ids"].shape
																						# Not using the same length for past_key_values
																						_UpperCamelCase			:			Optional[int]   			=						seqlen + 2
																						_UpperCamelCase			:			str   			=						(
																						    batch,
																						    self.num_attention_heads,
																						    past_key_values_length,
																						    self._config.hidden_size // self.num_attention_heads,
																						)
																						_UpperCamelCase			:			Dict   			=						[
																						    (torch.zeros(lowercase__					), torch.zeros(lowercase__					)) for _ in range(self.num_layers					)
																						]

										_UpperCamelCase			:			str   			=						common_inputs["attention_mask"]
										if self.use_past:
																_UpperCamelCase			:			int   			=						ordered_inputs["attention_mask"].dtype
																_UpperCamelCase			:			Optional[int]   			=						torch.cat(
																    [ordered_inputs["attention_mask"], torch.ones(lowercase__       ,		lowercase__       ,		dtype=lowercase__					)]       ,		dim=1					)

										return ordered_inputs







				@property
				def  snake_case__				(						self					:						Tuple					)							->int:


										'''simple docstring'''




										return 13


