






"""simple docstring"""





from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


__lowerCamelCase							=  logging.get_logger(__name__)

__lowerCamelCase							=  {
    'google/umt5-small': 'https://huggingface.co/google/umt5-small/resolve/main/config.json',
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}



class       __A		(   SCREAMING_SNAKE_CASE_			):
					UpperCAmelCase__								=				"umt5"
					UpperCAmelCase__								=				["past_key_values"]

					def __init__(			self  :      int     ,		__snake_case  :      Dict=2_5_0_1_1_2     ,		__snake_case  :      Dict=5_1_2     ,		__snake_case  :      List[Any]=6_4     ,		__snake_case  :      Union[str, Any]=1_0_2_4     ,		__snake_case  :      Optional[Any]=8     ,		__snake_case  :      Union[str, Any]=None     ,		__snake_case  :      Dict=6     ,		__snake_case  :      Any=3_2     ,		__snake_case  :      Any=1_2_8     ,		__snake_case  :      Optional[Any]=0.1     ,		__snake_case  :      Optional[int]=1E-6     ,		__snake_case  :      List[Any]=1.0     ,		__snake_case  :      Tuple="gated-gelu"     ,		__snake_case  :      Any=True     ,		__snake_case  :      Any=True     ,		__snake_case  :      Union[str, Any]="T5Tokenizer"     ,		__snake_case  :      Dict=True     ,		__snake_case  :      Union[str, Any]=0     ,		__snake_case  :      int=1     ,		__snake_case  :      Optional[Any]=0     ,		**__snake_case  :      Optional[int]     ,		)			->							str:
											super().__init__(
											    is_encoder_decoder=__snake_case     ,		tokenizer_class=__snake_case     ,		tie_word_embeddings=__snake_case     ,		pad_token_id=__snake_case     ,		eos_token_id=__snake_case     ,		decoder_start_token_id=__snake_case     ,		**__snake_case     ,		)
											__magic_name__:	List[Any]							=       vocab_size
											__magic_name__:	Dict							=       d_model
											__magic_name__:	List[Any]							=       d_kv
											__magic_name__:	Any							=       d_ff
											__magic_name__:	List[str]							=       num_layers
											__magic_name__:	Any							=       (
											    num_decoder_layers if num_decoder_layers is not None else self.num_layers
											)  # default = symmetry
											__magic_name__:	int							=       num_heads
											__magic_name__:	List[Any]							=       relative_attention_num_buckets
											__magic_name__:	List[str]							=       relative_attention_max_distance
											__magic_name__:	str							=       dropout_rate
											__magic_name__:	Optional[Any]							=       layer_norm_epsilon
											__magic_name__:	List[Any]							=       initializer_factor
											__magic_name__:	Union[str, Any]							=       feed_forward_proj
											__magic_name__:	int							=       use_cache

											__magic_name__:	int							=       self.feed_forward_proj.split("""-"""					)
											__magic_name__:	Tuple							=       act_info[-1]
											__magic_name__:	Dict							=       act_info[0] == """gated"""

											if len(__snake_case					) > 1 and act_info[0] != "gated" or len(__snake_case					) > 2:
																	raise ValueError(
																	    F'`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.'
																	    """Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. """
																	    """'gated-gelu' or 'relu'"""					)

											if feed_forward_proj == "gated-gelu":
																	__magic_name__:	Union[str, Any]							=       """gelu_new"""

					@property
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Dict:
											return self.d_model

					@property
					def   lowerCamelCase__   (			self  :      List[str]					)			->							Any:
											return self.num_heads






					@property
					def   lowerCamelCase__   (			self  :      Tuple					)			->							str:
											return self.num_layers



class       __A		(   SCREAMING_SNAKE_CASE_			):

					@property
					# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
					def   lowerCamelCase__   (			self  :      List[str]					)			->							Mapping[str, Mapping[int, str]]:
											__magic_name__:	Optional[int]							=       {
											    """input_ids""": {0: """batch""", 1: """encoder_sequence"""},
											    """attention_mask""": {0: """batch""", 1: """encoder_sequence"""},
											}
											if self.use_past:
																	__magic_name__:	str							=       """past_encoder_sequence + sequence"""
																	__magic_name__:	Any							=       {0: """batch"""}
																	__magic_name__:	List[Any]							=       {0: """batch""", 1: """past_decoder_sequence + sequence"""}
											else:
																	__magic_name__:	Union[str, Any]							=       {0: """batch""", 1: """decoder_sequence"""}
																	__magic_name__:	Tuple							=       {0: """batch""", 1: """decoder_sequence"""}

											if self.use_past:
																	self.fill_with_past_key_values_(__snake_case     ,		direction="""inputs"""					)

											return common_inputs

					@property
					# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							int:
											return 1_3






					@property
					def   lowerCamelCase__   (			self  :      Dict					)			->							float:
											return 5E-4











"""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							=  None
__lowerCamelCase							=  {
    '7B': 1_10_08,
    '13B': 1_38_24,
    '30B': 1_79_20,
    '65B': 2_20_16,
    '70B': 2_86_72,
}
__lowerCamelCase							=  {
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}



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



def   a    (     __UpperCAmelCase      :    Union[str, Any]  )    ->							Optional[Any]:
						with open(__UpperCAmelCase      ,   """r"""  ) as f:
												return json.load(__UpperCAmelCase  )



def   a    (     __UpperCAmelCase      :    Tuple      ,   __UpperCAmelCase      :    Union[str, Any]  )    ->							Optional[int]:
						with open(__UpperCAmelCase      ,   """w"""  ) as f:
												json.dump(__UpperCAmelCase      ,   __UpperCAmelCase  )



def   a    (     __UpperCAmelCase      :    List[Any]      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    str      ,   __UpperCAmelCase      :    Tuple=True  )    ->							Tuple:
						os.makedirs(__UpperCAmelCase      ,   exist_ok=__UpperCAmelCase  )
						__magic_name__:	Any							=       os.path.join(__UpperCAmelCase      ,   """tmp"""  )
						os.makedirs(__UpperCAmelCase      ,   exist_ok=__UpperCAmelCase  )

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

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

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

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

																		__magic_name__:	Dict							=       {
																		    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(),
																		}
																		__magic_name__:	Union[str, Any]							=       permute(
																		    torch.cat(
																		        [
																		            loaded[i][f'layers.{layer_i}.attention.wq.weight'].view(__UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase  )
																		            for i in range(__UpperCAmelCase  )
																		        ]      ,   dim=0      ,   ).reshape(__UpperCAmelCase      ,   __UpperCAmelCase  )  )
																		__magic_name__:	int							=       permute(
																		    torch.cat(
																		        [
																		            loaded[i][f'layers.{layer_i}.attention.wk.weight'].view(
																		                __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase  )
																		            for i in range(__UpperCAmelCase  )
																		        ]      ,   dim=0      ,   ).reshape(__UpperCAmelCase      ,   __UpperCAmelCase  )      ,   __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase      ,   )
																		__magic_name__:	Dict							=       torch.cat(
																		    [
																		        loaded[i][f'layers.{layer_i}.attention.wv.weight'].view(
																		            __UpperCAmelCase      ,   __UpperCAmelCase      ,   __UpperCAmelCase  )
																		        for i in range(__UpperCAmelCase  )
																		    ]      ,   dim=0      ,   ).reshape(__UpperCAmelCase      ,   __UpperCAmelCase  )

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

												__magic_name__:	Union[str, Any]							=       inv_freq
												for k, v in state_dict.items():
																		__magic_name__:	Tuple							=       filename
																		param_count += v.numel()
												torch.save(__UpperCAmelCase      ,   os.path.join(__UpperCAmelCase      ,   __UpperCAmelCase  )  )

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

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

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

						# 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."""  )
						__magic_name__:	Dict							=       LlamaForCausalLM.from_pretrained(__UpperCAmelCase      ,   torch_dtype=torch.floataa      ,   low_cpu_mem_usage=__UpperCAmelCase  )
						# Avoid saving this as part of the config.
						del model.config._name_or_path

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



def   a    (     __UpperCAmelCase      :    Dict      ,   __UpperCAmelCase      :    Dict  )    ->							List[Any]:
						# Initialize the tokenizer based on the `spm` model
						__magic_name__:	Optional[int]							=       LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
						print(f'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'  )
						__magic_name__:	Dict							=       tokenizer_class(__UpperCAmelCase  )
						tokenizer.save_pretrained(__UpperCAmelCase  )



def   a    (     )    ->							Any:
						__magic_name__:	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=__UpperCAmelCase      ,   help="""Whether or not to save using `safetensors`."""  )
						__magic_name__:	str							=       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      ,   )
						__magic_name__:	List[Any]							=       os.path.join(args.input_dir      ,   """tokenizer.model"""  )
						write_tokenizer(args.output_dir      ,   __UpperCAmelCase  )


if __name__ == "__main__":
							main()



