



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""": 11008,
    """13B""": 13824,
    """30B""": 17920,
    """65B""": 22016,
    """70B""": 28672,
}
_lowerCamelCase						    =  {
    """7B""": 1,
    """7Bf""": 1,
    """13B""": 2,
    """13Bf""": 2,
    """30B""": 4,
    """65B""": 8,
    """70B""": 8,
    """70Bf""": 8,
}






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     str=1		,	__lowerCamelCase   :     List[str]=256     ):



 """simple docstring"""





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






def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 with open(a_		,	"r"     ) as f:
  return json.load(a_     )






def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[Any]		,	__lowerCamelCase   :     List[Any]     ):



 """simple docstring"""





 with open(a_		,	"w"     ) as f:
  json.dump(a_		,	a_     )






def 	_lowerCAmelCase						(				__lowerCamelCase   :     int		,	__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     List[Any]		,	__lowerCamelCase   :     Any=True     ):



 """simple docstring"""





 os.makedirs(a_		,	exist_ok=a_     )
 __SCREAMING_SNAKE_CASE       :    Any       =	os.path.join(a_		,	"tmp"     )
 os.makedirs(a_		,	exist_ok=a_     )

 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	read_json(os.path.join(a_		,	"params.json"     )     )
 __SCREAMING_SNAKE_CASE       :    List[str]       =	NUM_SHARDS[model_size]
 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	params['''n_layers''']
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	params['''n_heads''']
 __SCREAMING_SNAKE_CASE       :    Any       =	n_heads // num_shards
 __SCREAMING_SNAKE_CASE       :    Tuple       =	params['''dim''']
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	dim // n_heads
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	10000.0
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	1.0 / (base ** (torch.arange(0		,	a_		,	2     ).float() / dims_per_head))

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

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

 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       :    Any       =	torch.load(os.path.join(a_		,	"consolidated.00.pth"     )		,	map_location="cpu"     )
 else:
  # Sharded
  __SCREAMING_SNAKE_CASE       :    Dict       =	[
      torch.load(os.path.join(a_		,	F"""consolidated.{i:02d}.pth"""     )		,	map_location="cpu"     )
      for i in range(a_     )
  ]
 __SCREAMING_SNAKE_CASE       :    Tuple       =	0
 __SCREAMING_SNAKE_CASE       :    Any       =	{'''weight_map''': {}}
 for layer_i in range(a_     ):
  __SCREAMING_SNAKE_CASE       :    Optional[int]       =	F"""pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"""
  if model_size == "7B":
   # Unsharded
   __SCREAMING_SNAKE_CASE       :    Tuple       =	{
       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       :    int       =	{
       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       :    Union[str, Any]       =	permute(
       torch.cat(
           [
               loaded[i][F"""layers.{layer_i}.attention.wq.weight"""].view(a_		,	a_		,	a_     )
               for i in range(a_     )
           ]		,	dim=0		,	).reshape(a_		,	a_     )     )
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	permute(
       torch.cat(
           [
               loaded[i][F"""layers.{layer_i}.attention.wk.weight"""].view(
                   a_		,	a_		,	a_     )
               for i in range(a_     )
           ]		,	dim=0		,	).reshape(a_		,	a_     )		,	a_		,	a_		,	a_		,	)
   __SCREAMING_SNAKE_CASE       :    Dict       =	torch.cat(
       [
           loaded[i][F"""layers.{layer_i}.attention.wv.weight"""].view(
               a_		,	a_		,	a_     )
           for i in range(a_     )
       ]		,	dim=0		,	).reshape(a_		,	a_     )

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

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

 __SCREAMING_SNAKE_CASE       :    List[Any]       =	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       :    List[Any]       =	{
      '''model.norm.weight''': loaded[0]['''norm.weight'''],
      '''model.embed_tokens.weight''': torch.cat(
          [loaded[i]["tok_embeddings.weight"] for i in range(a_     )]		,	dim=1     ),
      '''lm_head.weight''': torch.cat([loaded[i]["output.weight"] for i in range(a_     )]		,	dim=0     ),
  }

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

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

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

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






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any		,	__lowerCamelCase   :     str     ):



 """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       :    str       =	tokenizer_class(a_     )
 tokenizer.save_pretrained(a_     )






def 	_lowerCAmelCase						(				):



 """simple docstring"""





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


if __name__ == "__main__":
      main()











from diffusers.utils.testing_utils import require_onnxruntime

@require_onnxruntime
class       _SCREAMING_SNAKE_CASE	:
     pass






