





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


def _A       (   __magic_name__  , __magic_name__=1  , __magic_name__=256	):
   return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3	)	) + multiple_of - 1) // multiple_of)


def _A       (   __magic_name__	):
   with open(__magic_name__  , "r"	) as f:
      return json.load(__magic_name__	)


def _A       (   __magic_name__  , __magic_name__	):
   with open(__magic_name__  , "w"	) as f:
      json.dump(__magic_name__  , __magic_name__	)


def _A       (   __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__=True	):
   os.makedirs(__magic_name__  , exist_ok=__magic_name__	)
   lowercase__			       =  os.path.join(__magic_name__  , "tmp"	)
   os.makedirs(__magic_name__  , exist_ok=__magic_name__	)

   lowercase__			       =  read_json(os.path.join(__magic_name__  , "params.json"	)	)
   lowercase__			       =  NUM_SHARDS[model_size]
   lowercase__			       =  params["n_layers"]
   lowercase__			       =  params["n_heads"]
   lowercase__			       =  n_heads // num_shards
   lowercase__			       =  params["dim"]
   lowercase__			       =  dim // n_heads
   lowercase__			       =  10_000.0
   lowercase__			       =  1.0 / (base ** (torch.arange(0  , __magic_name__  , 2	).float() / dims_per_head))

   if "n_kv_heads" in params:
      lowercase__			       =  params["n_kv_heads"]  # for GQA / MQA
      lowercase__			       =  n_heads_per_shard // num_key_value_heads
      lowercase__			       =  dim // num_key_value_heads
   else:  # compatibility with other checkpoints
      lowercase__			       =  n_heads
      lowercase__			       =  n_heads_per_shard
      lowercase__			       =  dim

   # permute for sliced rotary
   def permute(__magic_name__  , __magic_name__=n_heads  , __magic_name__=dim  , __magic_name__=dim	):
      return w.view(__magic_name__  , dima // n_heads // 2  , 2  , __magic_name__	).transpose(1  , 2	).reshape(__magic_name__  , __magic_name__	)

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

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

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

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

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

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

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

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

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


def _A       (   __magic_name__  , __magic_name__	):
   # Initialize the tokenizer based on the `spm` model
   lowercase__			       =  LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
   print(f'''Saving a {tokenizer_class.__name__} to {tokenizer_path}.'''	)
   lowercase__			       =  tokenizer_class(__magic_name__	)
   tokenizer.save_pretrained(__magic_name__	)


def _A       (   ):
   lowercase__			       =  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=__magic_name__  , help="Whether or not to save using `safetensors`."	)
   lowercase__			       =  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  , )
   lowercase__			       =  os.path.join(args.input_dir  , "tokenizer.model"	)
   write_tokenizer(args.output_dir  , __magic_name__	)


if __name__ == "__main__":
  main()













from __future__ import annotations


def _A       (   __magic_name__  , __magic_name__	):
   lowercase__			       =  []
   create_all_state(1  , __magic_name__  , __magic_name__  , []  , __magic_name__	)
   return result


def _A       (   __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , ):
   if level == 0:
      total_list.append(current_list[:]	)
      return

   for i in range(__magic_name__  , total_number - level + 2	):
      current_list.append(__magic_name__	)
      create_all_state(i + 1  , __magic_name__  , level - 1  , __magic_name__  , __magic_name__	)
      current_list.pop()


def _A       (   __magic_name__	):
   for i in total_list:
      print(*__magic_name__	)


if __name__ == "__main__":
  _snake_case           =			4
  _snake_case           =			2
  _snake_case           =			generate_all_combinations(n, k)
  print_all_state(total_list)






