





"""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"
					)
					lowerCamelCase__ :       Any						     =			None
lowerCamelCase__ :       List[str]						     =			{
    '7B': 11008,
    '13B': 13824,
    '30B': 17920,
    '65B': 22016,
    '70B': 28672,
}
lowerCamelCase__ :       Tuple						     =			{
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}





def __A						(   a_ :							List[str]		,		a_ :							Optional[Any]=1		,		a_ :							List[str]=2_56     )->							str:




		'''simple docstring'''






		return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3     )     ) + multiple_of - 1) // multiple_of)





def __A						(   a_ :							Optional[int]     )->							Any:




		'''simple docstring'''






		with open(lowerCAmelCase_		,		'''r'''     ) as f:
				return json.load(lowerCAmelCase_     )





def __A						(   a_ :							Union[str, Any]		,		a_ :							Union[str, Any]     )->							List[Any]:




		'''simple docstring'''






		with open(lowerCAmelCase_		,		'''w'''     ) as f:
				json.dump(lowerCAmelCase_		,		lowerCAmelCase_     )





def __A						(   a_ :							Tuple		,		a_ :							int		,		a_ :							Tuple		,		a_ :							Optional[int]=True     )->							List[str]:




		'''simple docstring'''






		os.makedirs(lowerCAmelCase_		,		exist_ok=lowerCAmelCase_     )
		SCREAMING_SNAKE_CASE       :     Tuple				= os.path.join(lowerCAmelCase_		,		'''tmp'''     )
		os.makedirs(lowerCAmelCase_		,		exist_ok=lowerCAmelCase_     )

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

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

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

						SCREAMING_SNAKE_CASE       :     str				= {
						    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(),
						}
						SCREAMING_SNAKE_CASE       :     int				= 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_     )     )
						SCREAMING_SNAKE_CASE       :     Any				= 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_		,		)
						SCREAMING_SNAKE_CASE       :     str				= 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_     )

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

				SCREAMING_SNAKE_CASE       :     Dict				= inv_freq
				for k, v in state_dict.items():
						SCREAMING_SNAKE_CASE       :     List[str]				= filename
						param_count += v.numel()
				torch.save(lowerCAmelCase_		,		os.path.join(lowerCAmelCase_		,		lowerCAmelCase_     )     )

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

		# Write configs
		SCREAMING_SNAKE_CASE       :     Optional[Any]				= {'''total_size''': param_count * 2}
		write_json(lowerCAmelCase_		,		os.path.join(lowerCAmelCase_		,		'''pytorch_model.bin.index.json'''     )     )
		SCREAMING_SNAKE_CASE       :     Optional[int]				= params['''ffn_dim_multiplier'''] if '''ffn_dim_multiplier''' in params else 1
		SCREAMING_SNAKE_CASE       :     str				= params['''multiple_of'''] if '''multiple_of''' in params else 2_56
		SCREAMING_SNAKE_CASE       :     Optional[int]				= 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.'''     )
		SCREAMING_SNAKE_CASE       :     Optional[int]				= 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						(   a_ :							Union[str, Any]		,		a_ :							Union[str, Any]     )->							Tuple:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     Optional[int]				= LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
		print(F"Saving a {tokenizer_class.__name__} to {tokenizer_path}."     )
		SCREAMING_SNAKE_CASE       :     Optional[int]				= tokenizer_class(lowerCAmelCase_     )
		tokenizer.save_pretrained(lowerCAmelCase_     )





def __A						(   )->							Any:




		'''simple docstring'''






		SCREAMING_SNAKE_CASE       :     List[Any]				= 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`.'''     )
		SCREAMING_SNAKE_CASE       :     int				= 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		,		)
		SCREAMING_SNAKE_CASE       :     List[str]				= os.path.join(args.input_dir		,		'''tokenizer.model'''     )
		write_tokenizer(args.output_dir		,		lowerCAmelCase_     )


if __name__ == "__main__":
					main()







"""simple docstring"""






import json
import os
from typing import Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCamelCase__ :       List[str]						     =			logging.get_logger(__name__)

lowerCamelCase__ :       Optional[int]						     =			{"vocab_file": "vocab.json"}

lowerCamelCase__ :       Dict						     =			{
    "vocab_file": {
        "mgp-str": "https://huggingface.co/alibaba-damo/mgp-str-base/blob/main/vocab.json",
    }
}

lowerCamelCase__ :       Optional[Any]						     =			{"mgp-str": 27}




class 				lowercase__(				_UpperCAmelCase   ):
	'''simple docstring'''

	UpperCamelCase				=  VOCAB_FILES_NAMES
	UpperCamelCase				=  PRETRAINED_VOCAB_FILES_MAP
	UpperCamelCase				=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES





	def __init__(						self    :int		, lowerCamelCase_    :Optional[int]		, lowerCamelCase_    :Optional[int]="[GO]"		, lowerCamelCase_    :int="[GO]"		, lowerCamelCase_    :str="[s]"		, lowerCamelCase_    :Dict="[GO]"		, **lowerCamelCase_    :List[str]					)							->			Tuple:







			'''simple docstring'''




			super().__init__(
			    unk_token=lowerCamelCase_		, bos_token=lowerCamelCase_		, eos_token=lowerCamelCase_		, pad_token=lowerCamelCase_		, **lowerCamelCase_		, )

			with open(lowerCamelCase_		, encoding='''utf-8'''					) as vocab_handle:
					SCREAMING_SNAKE_CASE       :     int				= json.load(lowerCamelCase_					)
			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= {v: k for k, v in self.vocab.items()}





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







			'''simple docstring'''




			return len(self.vocab					)





	def   __lowerCAmelCase (						self    :List[str]					)							->			Dict:







			'''simple docstring'''




			return dict(self.vocab		, **self.added_tokens_encoder					)





	def   __lowerCAmelCase (						self    :List[str]		, lowerCamelCase_    :Optional[Any]					)							->			List[str]:







			'''simple docstring'''




			SCREAMING_SNAKE_CASE       :     Union[str, Any]				= []
			for s in text:
					char_tokens.extend(lowerCamelCase_					)
			return char_tokens





	def   __lowerCAmelCase (						self    :Dict		, lowerCamelCase_    :Tuple					)							->			Optional[int]:







			'''simple docstring'''




			return self.vocab.get(lowerCamelCase_		, self.vocab.get(self.unk_token					)					)





	def   __lowerCAmelCase (						self    :List[str]		, lowerCamelCase_    :Dict					)							->			Optional[int]:







			'''simple docstring'''




			return self.decoder.get(lowerCamelCase_					)







	def   __lowerCAmelCase (						self    :Optional[int]		, lowerCamelCase_    :str		, lowerCamelCase_    :Optional[str] = None					)							->			Tuple[str]:







			'''simple docstring'''




			if not os.path.isdir(lowerCamelCase_					):
					logger.error('''Vocabulary path ({}) should be a directory'''.format(lowerCamelCase_					)					)
					return
			SCREAMING_SNAKE_CASE       :     List[Any]				= os.path.join(
			    lowerCamelCase_		, (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']					)

			with open(lowerCamelCase_		, '''w'''		, encoding='''utf-8'''					) as f:
					f.write(json.dumps(self.vocab		, indent=2		, sort_keys=lowerCamelCase_		, ensure_ascii=lowerCamelCase_					) + '''\n'''					)

			return (vocab_file,)





