






'''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 	__UpperCAmelCase       ( __magic_name__						)->    List[Any]:



			"""simple docstring"""



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

			if "tiny" in model_name or "small" in model_name or "base" in model_name:
						snake_case_				:					str   =       [3, 3, 3, 3]
						if "lrf" in model_name:
									snake_case_				:					List[Any]   =       [3, 3, 3, 3]
						else:
									snake_case_				:					List[Any]   =       [2, 2, 2, 2]

			if "tiny" in model_name:
						snake_case_				:					List[Any]   =       96
			elif "small" in model_name:
						snake_case_				:					Any   =       96
			elif "base" in model_name:
						snake_case_				:					str   =       128
			elif "large" in model_name:
						snake_case_				:					List[Any]   =       192
			elif "xlarge" in model_name:
						snake_case_				:					int   =       256
			elif "huge" in model_name:
						snake_case_				:					str   =       352

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

			snake_case_				:					Union[str, Any]   =       json.load(open(hf_hub_download(__magic_name__     ,__magic_name__     ,repo_type="dataset"						)     ,"r"						)						)
			snake_case_				:					str   =       {int(__magic_name__						): v for k, v in idalabel.items()}
			snake_case_				:					str   =       {v: k for k, v in idalabel.items()}

			snake_case_				:					int   =       FocalNetConfig(
			    embed_dim=__magic_name__     ,depths=__magic_name__     ,focal_levels=__magic_name__     ,focal_windows=__magic_name__     ,use_conv_embed=__magic_name__     ,idalabel=__magic_name__     ,labelaid=__magic_name__     ,use_post_layernorm=__magic_name__     ,use_layerscale=__magic_name__     ,)

			return config

def 	__UpperCAmelCase       ( __magic_name__						)->    List[str]:



			"""simple docstring"""



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

			if name == "norm.weight":
						snake_case_				:					Any   =       "layernorm.weight"
			if name == "norm.bias":
						snake_case_				:					Any   =       "layernorm.bias"

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

			return name

def 	__UpperCAmelCase       ( __magic_name__     ,__magic_name__     ,__magic_name__=False						)->    Dict:



			"""simple docstring"""



			snake_case_				:					Tuple   =       {
			    "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

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

			# rename keys
			for key in state_dict.copy().keys():
						snake_case_				:					Optional[int]   =       state_dict.pop(__magic_name__						)
						snake_case_				:					List[str]   =       val

			snake_case_				:					str   =       get_focalnet_config(__magic_name__						)
			snake_case_				:					List[Any]   =       FocalNetForImageClassification(__magic_name__						)
			model.eval()

			# load state dict
			model.load_state_dict(__magic_name__						)

			# verify conversion
			snake_case_				:					Tuple   =       "http://images.cocodataset.org/val2017/000000039769.jpg"

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

			snake_case_				:					Union[str, Any]   =       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]						),
			    ]						)

			snake_case_				:					Optional[Any]   =       image_transforms(__magic_name__						).unsqueeze(0						)

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

			snake_case_				:					int   =       model(**__magic_name__						)

			snake_case_				:					Dict   =       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":
						snake_case_				:					Tuple   =       torch.tensor([0.2_166, -0.4_368, 0.2_191]						)
			elif model_name == "focalnet-tiny-lrf":
						snake_case_				:					Optional[Any]   =       torch.tensor([1.1_669, 0.0_125, -0.1_695]						)
			elif model_name == "focalnet-small":
						snake_case_				:					List[Any]   =       torch.tensor([0.4_917, -0.0_430, 0.1_341]						)
			elif model_name == "focalnet-small-lrf":
						snake_case_				:					Optional[int]   =       torch.tensor([-0.2_588, -0.5_342, -0.2_331]						)
			elif model_name == "focalnet-base":
						snake_case_				:					List[Any]   =       torch.tensor([-0.1_655, -0.4_090, -0.1_730]						)
			elif model_name == "focalnet-base-lrf":
						snake_case_				:					List[str]   =       torch.tensor([0.5_306, -0.0_483, -0.3_928]						)
			assert torch.allclose(outputs.logits[0, :3]     ,__magic_name__     ,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(__magic_name__						)
						processor.save_pretrained(__magic_name__						)

			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   :					Optional[int]   							=					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   :					List[Any]   							=					parser.parse_args()
		convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












'''simple docstring'''




from typing import TYPE_CHECKING

from ...file_utils import _LazyModule, is_torch_available
from ...utils import OptionalDependencyNotAvailable


__lowerCamelCase   :					Dict   							=					{
    '''configuration_gpt_neox_japanese''': ['''GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GPTNeoXJapaneseConfig'''],
    '''tokenization_gpt_neox_japanese''': ['''GPTNeoXJapaneseTokenizer'''],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCamelCase   :					int   							=					[
		    '''GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''GPTNeoXJapaneseForCausalLM''',
		    '''GPTNeoXJapaneseLayer''',
		    '''GPTNeoXJapaneseModel''',
		    '''GPTNeoXJapanesePreTrainedModel''',
		]


if TYPE_CHECKING:
		from .configuration_gpt_neox_japanese import GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXJapaneseConfig
		from .tokenization_gpt_neox_japanese import GPTNeoXJapaneseTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_gpt_neox_japanese import (
				    GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST,
				    GPTNeoXJapaneseForCausalLM,
				    GPTNeoXJapaneseLayer,
				    GPTNeoXJapaneseModel,
				    GPTNeoXJapanesePreTrainedModel,
				)


else:
		import sys

		__lowerCamelCase   :					Optional[int]   							=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




