




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"
		)
		__UpperCamelCase					:       List[Any]    					=      None
__UpperCamelCase					:       Dict    					=      {
    "7B": 11_008,
    "13B": 13_824,
    "30B": 17_920,
    "65B": 22_016,
    "70B": 28_672,
}
__UpperCamelCase					:       int    					=      {
    "7B": 1,
    "7Bf": 1,
    "13B": 2,
    "13Bf": 2,
    "30B": 4,
    "65B": 8,
    "70B": 8,
    "70Bf": 8,
}





def 		__A		(	__lowerCamelCase							,			__lowerCamelCase=1							,			__lowerCamelCase=256    )      ->       Any:
	return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3    )    ) + multiple_of - 1) // multiple_of)





def 		__A		(	__lowerCamelCase    )      ->       Dict:
	with open(__lowerCamelCase							,			"""r"""    ) as f:
		return json.load(__lowerCamelCase    )





def 		__A		(	__lowerCamelCase							,			__lowerCamelCase    )      ->       Optional[int]:
	with open(__lowerCamelCase							,			"""w"""    ) as f:
		json.dump(__lowerCamelCase							,			__lowerCamelCase    )





def 		__A		(	__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase=True    )      ->       Optional[int]:
	os.makedirs(__lowerCamelCase							,			exist_ok=__lowerCamelCase    )
	a										=							os.path.join(__lowerCamelCase							,			"""tmp"""    )
	os.makedirs(__lowerCamelCase							,			exist_ok=__lowerCamelCase    )

	a										=							read_json(os.path.join(__lowerCamelCase							,			"""params.json"""    )    )
	a										=							NUM_SHARDS[model_size]
	a										=							params["""n_layers"""]
	a										=							params["""n_heads"""]
	a										=							n_heads // num_shards
	a										=							params["""dim"""]
	a										=							dim // n_heads
	a										=							10000.0
	a										=							1.0 / (base ** (torch.arange(0							,			__lowerCamelCase							,			2    ).float() / dims_per_head))

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

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

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

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

	a										=							f'pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin'
	if model_size == "7B":
		# Unsharded
		a										=							{
		    """model.embed_tokens.weight""": loaded["""tok_embeddings.weight"""],
		    """model.norm.weight""": loaded["""norm.weight"""],
		    """lm_head.weight""": loaded["""output.weight"""],
		}
	else:
		a										=							{
		    """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():
		a										=							filename
		param_count += v.numel()
	torch.save(__lowerCamelCase							,			os.path.join(__lowerCamelCase							,			__lowerCamelCase    )    )

	# Write configs
	a										=							{"""total_size""": param_count * 2}
	write_json(__lowerCamelCase							,			os.path.join(__lowerCamelCase							,			"""pytorch_model.bin.index.json"""    )    )
	a										=							params["""ffn_dim_multiplier"""] if """ffn_dim_multiplier""" in params else 1
	a										=							params["""multiple_of"""] if """multiple_of""" in params else 256
	a										=							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."""    )
	a										=							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 		__A		(	__lowerCamelCase							,			__lowerCamelCase    )      ->       Optional[int]:
	# Initialize the tokenizer based on the `spm` model
	a										=							LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
	print(f'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'    )
	a										=							tokenizer_class(__lowerCamelCase    )
	tokenizer.save_pretrained(__lowerCamelCase    )





def 		__A		(	)      ->       List[Any]:
	a										=							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`."""    )
	a										=							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										=							os.path.join(args.input_dir							,			"""tokenizer.model"""    )
	write_tokenizer(args.output_dir							,			__lowerCamelCase    )


if __name__ == "__main__":
		main()







import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp





def 		__A		(	)      ->       Any:
	a										=							ArgumentParser(
	    description=(
	        """PyTorch TPU distributed training launch """
	        """helper utility that will spawn up """
	        """multiple distributed processes"""
	    )    )

	# Optional arguments for the launch helper
	parser.add_argument("""--num_cores"""							,			type=__lowerCamelCase							,			default=1							,			help="""Number of TPU cores to use (1 or 8)."""    )

	# positional
	parser.add_argument(
	    """training_script"""							,			type=__lowerCamelCase							,			help=(
	        """The full path to the single TPU training """
	        """program/script to be launched in parallel, """
	        """followed by all the arguments for the """
	        """training script"""
	    )							,			)

	# rest from the training program
	parser.add_argument("""training_script_args"""							,			nargs=__lowerCamelCase    )

	return parser.parse_args()





def 		__A		(	)      ->       Union[str, Any]:
	a										=							parse_args()

	# Import training_script as a module.
	a										=							Path(args.training_script    )
	sys.path.append(str(script_fpath.parent.resolve()    )    )
	a										=							script_fpath.stem
	a										=							importlib.import_module(__lowerCamelCase    )

	# Patch sys.argv
	a										=							[args.training_script] + args.training_script_args + ["""--tpu_num_cores""", str(args.num_cores    )]
	xmp.spawn(mod._mp_fn							,			args=()							,			nprocs=args.num_cores    )


if __name__ == "__main__":
		main()

