





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'''
	)
	_SCREAMING_SNAKE_CASE   =      None
_SCREAMING_SNAKE_CASE   =      {
    '''7B''': 1_1_0_0_8,
    '''13B''': 1_3_8_2_4,
    '''30B''': 1_7_9_2_0,
    '''65B''': 2_2_0_1_6,
    '''70B''': 2_8_6_7_2,
}
_SCREAMING_SNAKE_CASE   =      {
    '''7B''': 1,
    '''7Bf''': 1,
    '''13B''': 2,
    '''13Bf''': 2,
    '''30B''': 4,
    '''65B''': 8,
    '''70B''': 8,
    '''70Bf''': 8,
}




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[int]						,    lowerCamelCase_     :					List[str]=1						,    lowerCamelCase_     :					List[str]=2_5_6			):
			return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3			)			) + multiple_of - 1) // multiple_of)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					str			):
			with open(lowerCamelCase_						,    '''r'''			) as f:
						return json.load(lowerCamelCase_			)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					List[Any]						,    lowerCamelCase_     :					List[Any]			):
			with open(lowerCamelCase_						,    '''w'''			) as f:
						json.dump(lowerCamelCase_						,    lowerCamelCase_			)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					int						,    lowerCamelCase_     :					Tuple						,    lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					Dict=True			):
			os.makedirs(lowerCamelCase_						,    exist_ok=lowerCamelCase_			)
			__lowercase			  =							os.path.join(lowerCamelCase_						,    '''tmp'''			)
			os.makedirs(lowerCamelCase_						,    exist_ok=lowerCamelCase_			)

			__lowercase			  =							read_json(os.path.join(lowerCamelCase_						,    '''params.json'''			)			)
			__lowercase			  =							NUM_SHARDS[model_size]
			__lowercase			  =							params['''n_layers''']
			__lowercase			  =							params['''n_heads''']
			__lowercase			  =							n_heads // num_shards
			__lowercase			  =							params['''dim''']
			__lowercase			  =							dim // n_heads
			__lowercase			  =							1_0_0_0_0.0
			__lowercase			  =							1.0 / (base ** (torch.arange(0						,    lowerCamelCase_						,    2			).float() / dims_per_head))

			if "n_kv_heads" in params:
						__lowercase			  =							params['''n_kv_heads''']  # for GQA / MQA
						__lowercase			  =							n_heads_per_shard // num_key_value_heads
						__lowercase			  =							dim // num_key_value_heads
			else:  # compatibility with other checkpoints
						__lowercase			  =							n_heads
						__lowercase			  =							n_heads_per_shard
						__lowercase			  =							dim

			# permute for sliced rotary
			def permute(lowerCamelCase_     :					Tuple						,    lowerCamelCase_     :					Dict=n_heads						,    lowerCamelCase_     :					Tuple=dim						,    lowerCamelCase_     :					Dict=dim			):
						return w.view(lowerCamelCase_						,    dima // n_heads // 2						,    2						,    lowerCamelCase_			).transpose(1						,    2			).reshape(lowerCamelCase_						,    lowerCamelCase_			)

			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.)
						__lowercase			  =							torch.load(os.path.join(lowerCamelCase_						,    '''consolidated.00.pth'''			)						,    map_location='''cpu'''			)
			else:
						# Sharded
						__lowercase			  =							[
						    torch.load(os.path.join(lowerCamelCase_						,    f"consolidated.{i:02d}.pth"			)						,    map_location='''cpu'''			)
						    for i in range(lowerCamelCase_			)
						]
			__lowercase			  =							0
			__lowercase			  =							{'''weight_map''': {}}
			for layer_i in range(lowerCamelCase_			):
						__lowercase			  =							f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
						if model_size == "7B":
									# Unsharded
									__lowercase			  =							{
									    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.

									__lowercase			  =							{
									    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(),
									}
									__lowercase			  =							permute(
									    torch.cat(
									        [
									            loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_			)
									            for i in range(lowerCamelCase_			)
									        ]						,    dim=0						,    ).reshape(lowerCamelCase_						,    lowerCamelCase_			)			)
									__lowercase			  =							permute(
									    torch.cat(
									        [
									            loaded[i][f"layers.{layer_i}.attention.wk.weight"].view(
									                lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_			)
									            for i in range(lowerCamelCase_			)
									        ]						,    dim=0						,    ).reshape(lowerCamelCase_						,    lowerCamelCase_			)						,    lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_						,    )
									__lowercase			  =							torch.cat(
									    [
									        loaded[i][f"layers.{layer_i}.attention.wv.weight"].view(
									            lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_			)
									        for i in range(lowerCamelCase_			)
									    ]						,    dim=0						,    ).reshape(lowerCamelCase_						,    lowerCamelCase_			)

									__lowercase			  =							torch.cat(
									    [loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(lowerCamelCase_			)]						,    dim=1			)
									__lowercase			  =							torch.cat(
									    [loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(lowerCamelCase_			)]						,    dim=0			)
									__lowercase			  =							torch.cat(
									    [loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(lowerCamelCase_			)]						,    dim=1			)
									__lowercase			  =							torch.cat(
									    [loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(lowerCamelCase_			)]						,    dim=0			)

						__lowercase			  =							inv_freq
						for k, v in state_dict.items():
									__lowercase			  =							filename
									param_count += v.numel()
						torch.save(lowerCamelCase_						,    os.path.join(lowerCamelCase_						,    lowerCamelCase_			)			)

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

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

			# Write configs
			__lowercase			  =							{'''total_size''': param_count * 2}
			write_json(lowerCamelCase_						,    os.path.join(lowerCamelCase_						,    '''pytorch_model.bin.index.json'''			)			)
			__lowercase			  =							params['''ffn_dim_multiplier'''] if '''ffn_dim_multiplier''' in params else 1
			__lowercase			  =							params['''multiple_of'''] if '''multiple_of''' in params else 2_5_6
			__lowercase			  =							LlamaConfig(
			    hidden_size=lowerCamelCase_						,    intermediate_size=compute_intermediate_size(lowerCamelCase_						,    lowerCamelCase_						,    lowerCamelCase_			)						,    num_attention_heads=params['''n_heads''']						,    num_hidden_layers=params['''n_layers''']						,    rms_norm_eps=params['''norm_eps''']						,    num_key_value_heads=lowerCamelCase_						,    )
			config.save_pretrained(lowerCamelCase_			)

			# 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.'''			)
			__lowercase			  =							LlamaForCausalLM.from_pretrained(lowerCamelCase_						,    torch_dtype=torch.floataa						,    low_cpu_mem_usage=lowerCamelCase_			)
			# Avoid saving this as part of the config.
			del model.config._name_or_path

			print('''Saving in the Transformers format.'''			)
			model.save_pretrained(lowerCamelCase_						,    safe_serialization=lowerCamelCase_			)
			shutil.rmtree(lowerCamelCase_			)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					List[Any]			):
			# Initialize the tokenizer based on the `spm` model
			__lowercase			  =							LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
			print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}."			)
			__lowercase			  =							tokenizer_class(lowerCamelCase_			)
			tokenizer.save_pretrained(lowerCamelCase_			)




def 				_lowerCAmelCase  (			):
			__lowercase			  =							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=lowerCamelCase_						,    help='''Whether or not to save using `safetensors`.'''			)
			__lowercase			  =							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						,    )
			__lowercase			  =							os.path.join(args.input_dir						,    '''tokenizer.model'''			)
			write_tokenizer(args.output_dir						,    lowerCamelCase_			)


if __name__ == "__main__":
	main()












'''simple docstring'''



from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_SCREAMING_SNAKE_CASE   =      {'''configuration_focalnet''': ['''FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''FocalNetConfig''']}


try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_SCREAMING_SNAKE_CASE   =      [
	    '''FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''FocalNetForImageClassification''',
	    '''FocalNetForMaskedImageModeling''',
	    '''FocalNetBackbone''',
	    '''FocalNetModel''',
	    '''FocalNetPreTrainedModel''',
	]

if TYPE_CHECKING:
	from .configuration_focalnet import FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FocalNetConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_focalnet import (
		    FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST,
		    FocalNetBackbone,
		    FocalNetForImageClassification,
		    FocalNetForMaskedImageModeling,
		    FocalNetModel,
		    FocalNetPreTrainedModel,
		)

else:
	import sys

	_SCREAMING_SNAKE_CASE   =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





