




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







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_=1 ,				lowercase_=256						)   ->       Optional[Any]:



       """simple docstring"""

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







def SCREAMING_SNAKE_CASE				(       lowercase_						)   ->       Tuple:



       """simple docstring"""

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







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_						)   ->       int:



       """simple docstring"""

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







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_ ,				lowercase_ ,				lowercase_=True						)   ->       Tuple:



       """simple docstring"""

       os.makedirs(lowercase_ ,				exist_ok=lowercase_						)
       A__       =    os.path.join(lowercase_ ,				'''tmp'''						)
       os.makedirs(lowercase_ ,				exist_ok=lowercase_						)

       A__       =    read_json(os.path.join(lowercase_ ,				'''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__       =    1_00_00.0
       A__       =    1.0 / (base ** (torch.arange(0 ,				lowercase_ ,				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(lowercase_ ,				lowercase_=n_heads ,				lowercase_=dim ,				lowercase_=dim						):
              return w.view(lowercase_ ,				dima // n_heads // 2 ,				2 ,				lowercase_						).transpose(1 ,				2						).reshape(lowercase_ ,				lowercase_						)

       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(lowercase_ ,				'''consolidated.00.pth'''						) ,				map_location='''cpu'''						)
       else:
              # Sharded
              A__       =    [
                  torch.load(os.path.join(lowercase_ ,				f"""consolidated.{i:02d}.pth"""						) ,				map_location='''cpu'''						)
                  for i in range(lowercase_						)
              ]
       A__       =    0
       A__       =    {'''weight_map''': {}}
       for layer_i in range(lowercase_						):
              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(lowercase_ ,				lowercase_ ,				lowercase_						)
                                 for i in range(lowercase_						)
                             ] ,				dim=0 ,				).reshape(lowercase_ ,				lowercase_						)						)
                     A__       =    permute(
                         torch.cat(
                             [
                                 loaded[i][f"""layers.{layer_i}.attention.wk.weight"""].view(
                                     lowercase_ ,				lowercase_ ,				lowercase_						)
                                 for i in range(lowercase_						)
                             ] ,				dim=0 ,				).reshape(lowercase_ ,				lowercase_						) ,				lowercase_ ,				lowercase_ ,				lowercase_ ,				)
                     A__       =    torch.cat(
                         [
                             loaded[i][f"""layers.{layer_i}.attention.wv.weight"""].view(
                                 lowercase_ ,				lowercase_ ,				lowercase_						)
                             for i in range(lowercase_						)
                         ] ,				dim=0 ,				).reshape(lowercase_ ,				lowercase_						)

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

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

       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(lowercase_						)] ,				dim=1						),
                  '''lm_head.weight''': torch.cat([loaded[i]['''output.weight'''] for i in range(lowercase_						)] ,				dim=0						),
              }

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

       # Write configs
       A__       =    {'''total_size''': param_count * 2}
       write_json(lowercase_ ,				os.path.join(lowercase_ ,				'''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=lowercase_ ,				intermediate_size=compute_intermediate_size(lowercase_ ,				lowercase_ ,				lowercase_						) ,				num_attention_heads=params['''n_heads'''] ,				num_hidden_layers=params['''n_layers'''] ,				rms_norm_eps=params['''norm_eps'''] ,				num_key_value_heads=lowercase_ ,				)
       config.save_pretrained(lowercase_						)

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

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







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_						)   ->       str:



       """simple docstring"""

       A__       =    LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
       print(f"""Saving a {tokenizer_class.__name__} to {tokenizer_path}."""						)
       A__       =    tokenizer_class(lowercase_						)
       tokenizer.save_pretrained(lowercase_						)







def SCREAMING_SNAKE_CASE				(       )   ->       Dict:



       """simple docstring"""

       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=lowercase_ ,				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 ,				lowercase_						)


if __name__ == "__main__":
  main()








from __future__ import annotations







def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_ ,				lowercase_						)   ->       tuple[float, list[float]]:



       """simple docstring"""

       A__       =    list(range(len(lowercase_						)						)						)
       A__       =    [v / w for v, w in zip(lowercase_ ,				lowercase_						)]
       index.sort(key=lambda lowercase_						: ratio[i] ,				reverse=lowercase_						)

       A__       =    0
       A__       =    [0] * len(lowercase_						)
       for i in index:
              if weight[i] <= capacity:
                     A__       =    1
                     max_value += value[i]
                     capacity -= weight[i]
              else:
                     A__       =    capacity / weight[i]
                     max_value += value[i] * capacity / weight[i]
                     break

       return max_value, fractions


if __name__ == "__main__":
  import doctest

  doctest.testmod()


