





"""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'
							)
							_UpperCamelCase					      =   None
_UpperCamelCase					      =   {
    '7B': 1_1008,
    '13B': 1_3824,
    '30B': 1_7920,
    '65B': 2_2016,
    '70B': 2_8672,
}
_UpperCamelCase					      =   {
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Optional[Any]     ,    SCREAMING_SNAKE_CASE							:				List[Any]=1     ,    SCREAMING_SNAKE_CASE							:				List[Any]=256 ):


							'''simple docstring'''







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



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Dict ):


							'''simple docstring'''







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



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Optional[int]     ,    SCREAMING_SNAKE_CASE							:				Union[str, Any] ):


							'''simple docstring'''







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



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Union[str, Any]     ,    SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				int     ,    SCREAMING_SNAKE_CASE							:				List[Any]=True ):


							'''simple docstring'''







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

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

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

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

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

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

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

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

							for k, v in state_dict.items():
														__lowerCamelCase   :   Optional[int]       =filename
														param_count += v.numel()
							torch.save(SCREAMING_SNAKE_CASE     ,    os.path.join(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE ) )

							# Write configs
							__lowerCamelCase   :   Optional[int]       ={'''total_size''': param_count * 2}
							write_json(SCREAMING_SNAKE_CASE     ,    os.path.join(SCREAMING_SNAKE_CASE     ,    '''pytorch_model.bin.index.json''' ) )
							__lowerCamelCase   :   Optional[int]       =params['''ffn_dim_multiplier'''] if '''ffn_dim_multiplier''' in params else 1
							__lowerCamelCase   :   List[Any]       =params['''multiple_of'''] if '''multiple_of''' in params else 256
							__lowerCamelCase   :   Optional[Any]       =LlamaConfig(
							    hidden_size=SCREAMING_SNAKE_CASE     ,    intermediate_size=compute_intermediate_size(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )     ,    num_attention_heads=params['''n_heads''']     ,    num_hidden_layers=params['''n_layers''']     ,    rms_norm_eps=params['''norm_eps''']     ,    num_key_value_heads=SCREAMING_SNAKE_CASE     ,    )
							config.save_pretrained(SCREAMING_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.''' )
							__lowerCamelCase   :   Dict       =LlamaForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE     ,    torch_dtype=torch.floataa     ,    low_cpu_mem_usage=SCREAMING_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(SCREAMING_SNAKE_CASE     ,    safe_serialization=SCREAMING_SNAKE_CASE )
							shutil.rmtree(SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Union[str, Any]     ,    SCREAMING_SNAKE_CASE							:				Dict ):


							'''simple docstring'''







							__lowerCamelCase   :   List[Any]       =LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
							print(F'Saving a {tokenizer_class.__name__} to {tokenizer_path}.' )
							__lowerCamelCase   :   Optional[Any]       =tokenizer_class(SCREAMING_SNAKE_CASE )
							tokenizer.save_pretrained(SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   ):


							'''simple docstring'''







							__lowerCamelCase   :   int       =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=SCREAMING_SNAKE_CASE     ,    help='''Whether or not to save using `safetensors`.''' )
							__lowerCamelCase   :   Optional[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     ,    )
							__lowerCamelCase   :   Any       =os.path.join(args.input_dir     ,    '''tokenizer.model''' )
							write_tokenizer(args.output_dir     ,    SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
							main()











"""simple docstring"""





import argparse

from torch import nn

# transformers_old should correspond to branch `save_old_prophetnet_model_structure` here
# original prophetnet_checkpoints are saved under `patrickvonplaten/..._old` respectively
from transformers_old.modeling_prophetnet import (
    ProphetNetForConditionalGeneration as ProphetNetForConditionalGenerationOld,
)
from transformers_old.modeling_xlm_prophetnet import (
    XLMProphetNetForConditionalGeneration as XLMProphetNetForConditionalGenerationOld,
)

from transformers import ProphetNetForConditionalGeneration, XLMProphetNetForConditionalGeneration, logging


_UpperCamelCase					      =   logging.get_logger(__name__)
logging.set_verbosity_info()



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				str ):


							'''simple docstring'''







							if "xprophetnet" in prophetnet_checkpoint_path:
														__lowerCamelCase   :   Optional[Any]       =XLMProphetNetForConditionalGenerationOld.from_pretrained(SCREAMING_SNAKE_CASE )
														__lowerCamelCase     ,    __lowerCamelCase   :   List[Any]       =XLMProphetNetForConditionalGeneration.from_pretrained(
														    SCREAMING_SNAKE_CASE     ,    output_loading_info=SCREAMING_SNAKE_CASE )
							else:
														__lowerCamelCase   :   int       =ProphetNetForConditionalGenerationOld.from_pretrained(SCREAMING_SNAKE_CASE )
														__lowerCamelCase     ,    __lowerCamelCase   :   int       =ProphetNetForConditionalGeneration.from_pretrained(
														    SCREAMING_SNAKE_CASE     ,    output_loading_info=SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   Optional[Any]       =['''key_proj''', '''value_proj''', '''query_proj''']

							__lowerCamelCase   :   Tuple       ={
							    '''self_attn''': '''ngram_self_attn''',
							    '''cross_attn''': '''encoder_attn''',
							    '''cross_attn_layer_norm''': '''encoder_attn_layer_norm''',
							    '''feed_forward_layer_norm''': '''final_layer_norm''',
							    '''feed_forward''': '''''',
							    '''intermediate''': '''fc1''',
							    '''output''': '''fc2''',
							    '''key_proj''': '''k_proj''',
							    '''query_proj''': '''q_proj''',
							    '''value_proj''': '''v_proj''',
							    '''word_embeddings''': '''embed_tokens''',
							    '''embeddings_layer_norm''': '''emb_layer_norm''',
							    '''relative_pos_embeddings''': '''relative_linear''',
							    '''ngram_embeddings''': '''ngram_input_embed''',
							    '''position_embeddings''': '''embed_positions''',
							}

							for key in loading_info["missing_keys"]:
														__lowerCamelCase   :   int       =key.split('''.''' )

														if attributes[0] == "lm_head":
																					__lowerCamelCase   :   int       =prophet
																					__lowerCamelCase   :   Optional[int]       =prophet_old
														else:
																					__lowerCamelCase   :   Any       =prophet.prophetnet
																					__lowerCamelCase   :   Union[str, Any]       =prophet_old.model

														__lowerCamelCase   :   Optional[Any]       =False
														for attribute in attributes:
																					if attribute in mapping:
																												__lowerCamelCase   :   Optional[Any]       =mapping[attribute]
																												if not hasattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE ) and len(SCREAMING_SNAKE_CASE ) > 0:
																																			__lowerCamelCase   :   Any       =attribute
																					elif hasattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE ):
																												__lowerCamelCase   :   Any       =attribute

																					if attribute == "weight":
																												assert old_model.weight.shape == model.weight.shape, "Shapes have to match!"
																												__lowerCamelCase   :   str       =old_model.weight
																												logger.info(F'{attribute} is initialized.' )
																												__lowerCamelCase   :   Any       =True
																												break
																					elif attribute == "bias":
																												assert old_model.bias.shape == model.bias.shape, "Shapes have to match!"
																												__lowerCamelCase   :   Union[str, Any]       =old_model.bias
																												logger.info(F'{attribute} is initialized' )
																												__lowerCamelCase   :   str       =True
																												break
																					elif attribute in special_keys and hasattr(SCREAMING_SNAKE_CASE     ,    '''in_proj_weight''' ):
																												__lowerCamelCase   :   int       =old_model.in_proj_weight.shape[0] // 3
																												__lowerCamelCase   :   Union[str, Any]       =getattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )
																												param.weight.shape == old_model.in_proj_weight[:embed_dim, :].shape, "Shapes have to match"
																												param.bias.shape == old_model.in_proj_bias[:embed_dim].shape, "Shapes have to match"
																												if attribute == "query_proj":
																																			__lowerCamelCase   :   List[str]       =nn.Parameter(old_model.in_proj_weight[:embed_dim, :] )
																																			__lowerCamelCase   :   str       =nn.Parameter(old_model.in_proj_bias[:embed_dim] )

																												elif attribute == "key_proj":
																																			__lowerCamelCase   :   List[str]       =nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :] )
																																			__lowerCamelCase   :   Tuple       =nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim] )
																												elif attribute == "value_proj":
																																			__lowerCamelCase   :   Optional[Any]       =nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :] )
																																			__lowerCamelCase   :   int       =nn.Parameter(old_model.in_proj_bias[2 * embed_dim :] )
																												__lowerCamelCase   :   Dict       =True
																												break
																					elif attribute == "position_embeddings":
																												assert (
																												    model.position_embeddings.weight.shape[-1] == old_model.embed_positions.weight.shape[-1]
																												), "Hidden size has to match"
																												assert model.position_embeddings.weight.shape[0] == 512, "We want 512 position_embeddings."
																												__lowerCamelCase   :   str       =nn.Parameter(old_model.embed_positions.weight[:512, :] )
																												__lowerCamelCase   :   Dict       =True
																												break

																					if attribute.isdigit():
																												__lowerCamelCase   :   List[str]       =model[int(SCREAMING_SNAKE_CASE )]
																												__lowerCamelCase   :   Optional[Any]       =old_model[int(SCREAMING_SNAKE_CASE )]
																					else:
																												__lowerCamelCase   :   int       =getattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

																												if old_attribute == "":
																																			__lowerCamelCase   :   Dict       =old_model
																												else:
																																			if not hasattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE ):
																																										raise ValueError(F'{old_model} does not have {old_attribute}' )
																																			__lowerCamelCase   :   Tuple       =getattr(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

														if not is_key_init:
																					raise ValueError(F'{key} was not correctly initialized!' )

							print(F'Saving model to {pytorch_dump_folder_path}' )
							prophet.save_pretrained(SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
							_UpperCamelCase					      =   argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--prophetnet_checkpoint_path', default=None, type=str, required=True, help='Path the official PyTorch dump.'
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
							)
							_UpperCamelCase					      =   parser.parse_args()
							convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path)




