




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 				_A  (		lowerCAmelCase_	:   int  ,							lowerCAmelCase_	:   Optional[Any]=1  ,							lowerCAmelCase_	:   Tuple=256     ):
     """simple docstring"""

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



def 				_A  (		lowerCAmelCase_	:   Dict     ):
     """simple docstring"""

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



def 				_A  (		lowerCAmelCase_	:   Tuple  ,							lowerCAmelCase_	:   int     ):
     """simple docstring"""

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



def 				_A  (		lowerCAmelCase_	:   List[str]  ,							lowerCAmelCase_	:   str  ,							lowerCAmelCase_	:   Optional[Any]  ,							lowerCAmelCase_	:   List[str]=True     ):
     """simple docstring"""

     os.makedirs(lowerCAmelCase_  ,							exist_ok=lowerCAmelCase_     )
     lowerCAmelCase__            =					os.path.join(lowerCAmelCase_  ,							"tmp"     )
     os.makedirs(lowerCAmelCase_  ,							exist_ok=lowerCAmelCase_     )

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

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

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

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

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

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

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

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

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

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

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



def 				_A  (		lowerCAmelCase_	:   Tuple  ,							lowerCAmelCase_	:   int     ):
     """simple docstring"""

     lowerCAmelCase__            =					LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
     print(F'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'     )
     lowerCAmelCase__            =					tokenizer_class(lowerCAmelCase_     )
     tokenizer.save_pretrained(lowerCAmelCase_     )



def 				_A  (		):
     """simple docstring"""

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


if __name__ == "__main__":
     main()





'''simple docstring'''



from math import pow
def        UpperCAmelCase (    A	:					int       ,		A	:					int       ,		A	:					int       ,		A	:					int       ,		A	:					int       ,		):
							if current_sum == needed_sum:
														# If the sum of the powers is equal to needed_sum, then we have a solution.
														solutions_count += 1
														return current_sum, solutions_count

							SCREAMING_SNAKE_CASE       :   Tuple         =      int(pow(A       ,		A					)					)
							if current_sum + i_to_n <= needed_sum:
														# If the sum of the powers is less than needed_sum, then continue adding powers.
														current_sum += i_to_n
														SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Dict         =      backtrack(
														    A       ,		A       ,		current_number + 1       ,		A       ,		A					)
														current_sum -= i_to_n
							if i_to_n < needed_sum:
														# If the power of i is less than needed_sum, then try with the next power.
														SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   int         =      backtrack(
														    A       ,		A       ,		current_number + 1       ,		A       ,		A					)
							return current_sum, solutions_count
def        UpperCAmelCase (    A	:					int       ,		A	:					int					):
							if not (1 <= needed_sum <= 1000 and 2 <= power <= 10):
														raise ValueError(
														    '''Invalid input\n'''
														    '''needed_sum must be between 1 and 1000, power between 2 and 10.'''					)

							return backtrack(A       ,		A       ,		1       ,		0       ,		0					)[1]  # Return the solutions_count


if __name__ == "__main__":
							import doctest

							doctest.testmod()




