





'''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	(       __snake_case					:		Optional[Any]	)	->							str:






						"""simple docstring"""





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

						if "tiny" in model_name or "small" in model_name or "base" in model_name:
												A__					:       List[Any]     =[3, 3, 3, 3]
												if "lrf" in model_name:
																		A__					:       Dict     =[3, 3, 3, 3]
												else:
																		A__					:       int     =[2, 2, 2, 2]

						if "tiny" in model_name:
												A__					:       Dict     =96
						elif "small" in model_name:
												A__					:       Dict     =96
						elif "base" in model_name:
												A__					:       Dict     =128
						elif "large" in model_name:
												A__					:       int     =192
						elif "xlarge" in model_name:
												A__					:       int     =256
						elif "huge" in model_name:
												A__					:       Tuple     =352

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

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

						A__					:       Union[str, Any]     =FocalNetConfig(
						    embed_dim=_snake_case,		depths=_snake_case,		focal_levels=_snake_case,		focal_windows=_snake_case,		use_conv_embed=_snake_case,		idalabel=_snake_case,		labelaid=_snake_case,		use_post_layernorm=_snake_case,		use_layerscale=_snake_case,		)

						return config
def 						__lowerCamelCase	(       __snake_case					:		Optional[Any]	)	->							List[str]:






						"""simple docstring"""





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

						if name == "norm.weight":
												A__					:       str     ='''layernorm.weight'''
						if name == "norm.bias":
												A__					:       int     ='''layernorm.bias'''

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

						return name
def 						__lowerCamelCase	(       __snake_case					:		Any,		__snake_case					:		Optional[Any],		__snake_case					:		str=False	)	->							Dict:






						"""simple docstring"""





						A__					:       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__					:       Union[str, Any]     =model_name_to_url[model_name]
						print("""Checkpoint URL: """,		_snake_case	)
						A__					:       Any     =torch.hub.load_state_dict_from_url(_snake_case,		map_location="""cpu"""	)['''model''']

						# rename keys
						for key in state_dict.copy().keys():
												A__					:       Dict     =state_dict.pop(_snake_case	)
												A__					:       Dict     =val

						A__					:       int     =get_focalnet_config(_snake_case	)
						A__					:       List[Any]     =FocalNetForImageClassification(_snake_case	)
						model.eval()

						# load state dict
						model.load_state_dict(_snake_case	)

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

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

						A__					:       List[str]     =transforms.Compose(
						    [
						        transforms.Resize(256	),
						        transforms.CenterCrop(224	),
						        transforms.ToTensor(),
						        transforms.Normalize(mean=[0.4_85, 0.4_56, 0.4_06],		std=[0.2_29, 0.2_24, 0.2_25]	),
						    ]	)

						A__					:       str     =image_transforms(_snake_case	).unsqueeze(0	)

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

						A__					:       Any     =model(**_snake_case	)

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

						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__":
		__snake_case			:     Optional[Any]										=      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.',
		)

		__snake_case			:     List[str]										=      parser.parse_args()
		convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




'''simple docstring'''



import argparse
import gc
import json
import os
import shutil
import warnings

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
		from transformers import LlamaTokenizerFast
except ImportError as e:
		warnings.warn(e)
		warnings.warn(
		    'The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion'
		)
		__snake_case			:     Optional[Any]										=      None
__snake_case			:     Optional[Any]										=      {
    '7B': 1_1008,
    '13B': 1_3824,
    '30B': 1_7920,
    '65B': 2_2016,
    '70B': 2_8672,
}
__snake_case			:     Union[str, Any]										=      {
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}
def 						__lowerCamelCase	(       __snake_case					:		Optional[Any],		__snake_case					:		str=1,		__snake_case					:		Tuple=256	)	->							str:






						"""simple docstring"""





						return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3	)	) + multiple_of - 1) // multiple_of)
def 						__lowerCamelCase	(       __snake_case					:		Tuple	)	->							Tuple:






						"""simple docstring"""





						with open(__snake_case,		"""r"""	) as f:
												return json.load(__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		Optional[int],		__snake_case					:		Tuple	)	->							Dict:






						"""simple docstring"""





						with open(__snake_case,		"""w"""	) as f:
												json.dump(__snake_case,		__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		List[Any],		__snake_case					:		Any,		__snake_case					:		Any,		__snake_case					:		Tuple=True	)	->							Any:






						"""simple docstring"""





						os.makedirs(__snake_case,		exist_ok=__snake_case	)
						A__					:       List[Any]     =os.path.join(__snake_case,		"""tmp"""	)
						os.makedirs(__snake_case,		exist_ok=__snake_case	)

						A__					:       Dict     =read_json(os.path.join(__snake_case,		"""params.json"""	)	)
						A__					:       Dict     =NUM_SHARDS[model_size]
						A__					:       List[str]     =params["""n_layers"""]
						A__					:       int     =params["""n_heads"""]
						A__					:       str     =n_heads // num_shards
						A__					:       Tuple     =params["""dim"""]
						A__					:       Union[str, Any]     =dim // n_heads
						A__					:       str     =1_00_00.0
						A__					:       Any     =1.0 / (base ** (torch.arange(0,		__snake_case,		2	).float() / dims_per_head))

						if "n_kv_heads" in params:
												A__					:       Optional[Any]     =params["""n_kv_heads"""]  # for GQA / MQA
												A__					:       int     =n_heads_per_shard // num_key_value_heads
												A__					:       int     =dim // num_key_value_heads
						else:  # compatibility with other checkpoints
												A__					:       List[Any]     =n_heads
												A__					:       List[str]     =n_heads_per_shard
												A__					:       Dict     =dim

						# permute for sliced rotary
						def permute(__snake_case					:		Tuple,		__snake_case					:		Optional[int]=n_heads,		__snake_case					:		int=dim,		__snake_case					:		Optional[Any]=dim	):
												return w.view(__snake_case,		dima // n_heads // 2,		2,		__snake_case	).transpose(1,		2	).reshape(__snake_case,		__snake_case	)

						print(f"Fetching all parameters from the checkpoint at {input_base_path}."	)
						# Load weights
						if model_size == "7B":
												# Not sharded
												# (The sharded implementation would also work, but this is simpler.)
												A__					:       List[str]     =torch.load(os.path.join(__snake_case,		"""consolidated.00.pth"""	),		map_location="""cpu"""	)
						else:
												# Sharded
												A__					:       Optional[Any]     =[
												    torch.load(os.path.join(__snake_case,		f"consolidated.{i:02d}.pth"	),		map_location="""cpu"""	)
												    for i in range(__snake_case	)
												]
						A__					:       Optional[Any]     =0
						A__					:       str     ={"""weight_map""": {}}
						for layer_i in range(__snake_case	):
												A__					:       Dict     =f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
												if model_size == "7B":
																		# Unsharded
																		A__					:       Dict     ={
																		    f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
																		        loaded[f"layers.{layer_i}.attention.wq.weight"]	),
																		    f"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
																		        loaded[f"layers.{layer_i}.attention.wk.weight"]	),
																		    f"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[f"layers.{layer_i}.attention.wv.weight"],
																		    f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"layers.{layer_i}.attention.wo.weight"],
																		    f"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w1.weight"],
																		    f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w2.weight"],
																		    f"model.layers.{layer_i}.mlp.up_proj.weight": loaded[f"layers.{layer_i}.feed_forward.w3.weight"],
																		    f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"layers.{layer_i}.attention_norm.weight"],
																		    f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[f"layers.{layer_i}.ffn_norm.weight"],
																		}
												else:
																		# Sharded
																		# Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
																		# the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
																		# redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.

																		A__					:       Any     ={
																		    f"model.layers.{layer_i}.input_layernorm.weight": loaded[0][
																		        f"layers.{layer_i}.attention_norm.weight"
																		    ].clone(),
																		    f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[0][
																		        f"layers.{layer_i}.ffn_norm.weight"
																		    ].clone(),
																		}
																		A__					:       Optional[Any]     =permute(
																		    torch.cat(
																		        [
																		            loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(__snake_case,		__snake_case,		__snake_case	)
																		            for i in range(__snake_case	)
																		        ],		dim=0,		).reshape(__snake_case,		__snake_case	)	)
																		A__					:       int     =permute(
																		    torch.cat(
																		        [
																		            loaded[i][f"layers.{layer_i}.attention.wk.weight"].view(
																		                __snake_case,		__snake_case,		__snake_case	)
																		            for i in range(__snake_case	)
																		        ],		dim=0,		).reshape(__snake_case,		__snake_case	),		__snake_case,		__snake_case,		__snake_case,		)
																		A__					:       int     =torch.cat(
																		    [
																		        loaded[i][f"layers.{layer_i}.attention.wv.weight"].view(
																		            __snake_case,		__snake_case,		__snake_case	)
																		        for i in range(__snake_case	)
																		    ],		dim=0,		).reshape(__snake_case,		__snake_case	)

																		A__					:       List[str]     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(__snake_case	)],		dim=1	)
																		A__					:       Optional[int]     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(__snake_case	)],		dim=0	)
																		A__					:       str     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(__snake_case	)],		dim=1	)
																		A__					:       List[str]     =torch.cat(
																		    [loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(__snake_case	)],		dim=0	)

												A__					:       List[Any]     =inv_freq
												for k, v in state_dict.items():
																		A__					:       Optional[Any]     =filename
																		param_count += v.numel()
												torch.save(__snake_case,		os.path.join(__snake_case,		__snake_case	)	)

						A__					:       Tuple     =f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
						if model_size == "7B":
												# Unsharded
												A__					:       Tuple     ={
												    """model.embed_tokens.weight""": loaded["""tok_embeddings.weight"""],
												    """model.norm.weight""": loaded["""norm.weight"""],
												    """lm_head.weight""": loaded["""output.weight"""],
												}
						else:
												A__					:       Any     ={
												    """model.norm.weight""": loaded[0]["""norm.weight"""],
												    """model.embed_tokens.weight""": torch.cat(
												        [loaded[i]["""tok_embeddings.weight"""] for i in range(__snake_case	)],		dim=1	),
												    """lm_head.weight""": torch.cat([loaded[i]["""output.weight"""] for i in range(__snake_case	)],		dim=0	),
												}

						for k, v in state_dict.items():
												A__					:       int     =filename
												param_count += v.numel()
						torch.save(__snake_case,		os.path.join(__snake_case,		__snake_case	)	)

						# Write configs
						A__					:       Union[str, Any]     ={"""total_size""": param_count * 2}
						write_json(__snake_case,		os.path.join(__snake_case,		"""pytorch_model.bin.index.json"""	)	)
						A__					:       Optional[Any]     =params["""ffn_dim_multiplier"""] if """ffn_dim_multiplier""" in params else 1
						A__					:       List[Any]     =params["""multiple_of"""] if """multiple_of""" in params else 256
						A__					:       int     =LlamaConfig(
						    hidden_size=__snake_case,		intermediate_size=compute_intermediate_size(__snake_case,		__snake_case,		__snake_case	),		num_attention_heads=params["""n_heads"""],		num_hidden_layers=params["""n_layers"""],		rms_norm_eps=params["""norm_eps"""],		num_key_value_heads=__snake_case,		)
						config.save_pretrained(__snake_case	)

						# Make space so we can load the model properly now.
						del state_dict
						del loaded
						gc.collect()

						print("""Loading the checkpoint in a Llama model."""	)
						A__					:       List[Any]     =LlamaForCausalLM.from_pretrained(__snake_case,		torch_dtype=torch.floataa,		low_cpu_mem_usage=__snake_case	)
						# Avoid saving this as part of the config.
						del model.config._name_or_path

						print("""Saving in the Transformers format."""	)
						model.save_pretrained(__snake_case,		safe_serialization=__snake_case	)
						shutil.rmtree(__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		Union[str, Any],		__snake_case					:		Dict	)	->							Tuple:






						"""simple docstring"""





						A__					:       List[Any]     =LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
						print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}."	)
						A__					:       List[str]     =tokenizer_class(__snake_case	)
						tokenizer.save_pretrained(__snake_case	)
def 						__lowerCamelCase	(       )	->							Any:






						"""simple docstring"""





						A__					:       List[str]     =argparse.ArgumentParser()
						parser.add_argument(
						    """--input_dir""",		help="""Location of LLaMA weights, which contains tokenizer.model and model folders""",		)
						parser.add_argument(
						    """--model_size""",		choices=["""7B""", """7Bf""", """13B""", """13Bf""", """30B""", """65B""", """70B""", """70Bf""", """tokenizer_only"""],		)
						parser.add_argument(
						    """--output_dir""",		help="""Location to write HF model and tokenizer""",		)
						parser.add_argument("""--safe_serialization""",		type=__snake_case,		help="""Whether or not to save using `safetensors`."""	)
						A__					:       Any     =parser.parse_args()
						if args.model_size != "tokenizer_only":
												write_model(
												    model_path=args.output_dir,		input_base_path=os.path.join(args.input_dir,		args.model_size	),		model_size=args.model_size,		safe_serialization=args.safe_serialization,		)
						A__					:       List[Any]     =os.path.join(args.input_dir,		"""tokenizer.model"""	)
						write_tokenizer(args.output_dir,		__snake_case	)


if __name__ == "__main__":
		main()

