




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







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







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







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







def 				lowerCAmelCase_	(     snake_case_,snake_case_,snake_case_,snake_case_=True     ):
      os.makedirs(snake_case_,exist_ok=snake_case_     )
      _A							:     Tuple			=						os.path.join(snake_case_,"""tmp"""     )
      os.makedirs(snake_case_,exist_ok=snake_case_     )

      _A							:     int			=						read_json(os.path.join(snake_case_,"""params.json"""     )     )
      _A							:     Any			=						NUM_SHARDS[model_size]
      _A							:     Dict			=						params["""n_layers"""]
      _A							:     Optional[Any]			=						params["""n_heads"""]
      _A							:     Union[str, Any]			=						n_heads // num_shards
      _A							:     Dict			=						params["""dim"""]
      _A							:     Optional[Any]			=						dim // n_heads
      _A							:     List[Any]			=						1_00_00.0
      _A							:     List[str]			=						1.0 / (base ** (torch.arange(0,snake_case_,2     ).float() / dims_per_head))

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

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

                  _A							:     str			=						torch.cat(
                      [loaded[i][f'''layers.{layer_i}.attention.wo.weight'''] for i in range(snake_case_     )],dim=1     )
                  _A							:     Tuple			=						torch.cat(
                      [loaded[i][f'''layers.{layer_i}.feed_forward.w1.weight'''] for i in range(snake_case_     )],dim=0     )
                  _A							:     Dict			=						torch.cat(
                      [loaded[i][f'''layers.{layer_i}.feed_forward.w2.weight'''] for i in range(snake_case_     )],dim=1     )
                  _A							:     Any			=						torch.cat(
                      [loaded[i][f'''layers.{layer_i}.feed_forward.w3.weight'''] for i in range(snake_case_     )],dim=0     )

            _A							:     Any			=						inv_freq
            for k, v in state_dict.items():
                  _A							:     Dict			=						filename
                  param_count += v.numel()
            torch.save(snake_case_,os.path.join(snake_case_,snake_case_     )     )

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

      for k, v in state_dict.items():
            _A							:     List[str]			=						filename
            param_count += v.numel()
      torch.save(snake_case_,os.path.join(snake_case_,snake_case_     )     )

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







def 				lowerCAmelCase_	(     snake_case_,snake_case_     ):
      # Initialize the tokenizer based on the `spm` model
      _A							:     Union[str, Any]			=						LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
      print(f'''Saving a {tokenizer_class.__name__} to {tokenizer_path}.'''     )
      _A							:     List[Any]			=						tokenizer_class(snake_case_     )
      tokenizer.save_pretrained(snake_case_     )







def 				lowerCAmelCase_	(     ):
      _A							:     Dict			=						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=snake_case_,help="""Whether or not to save using `safetensors`."""     )
      _A							:     Union[str, 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,)
      _A							:     int			=						os.path.join(args.input_dir,"""tokenizer.model"""     )
      write_tokenizer(args.output_dir,snake_case_     )


if __name__ == "__main__":
  main()





"""simple docstring"""



def   __lowerCAmelCase						(						lowercase  :  int )	->						bool:


    """simple docstring"""





    if p < 2:
        raise ValueError("p should not be less than 2!" )
    elif p == 2:
        return True

    snake_case    :      Dict      					=    4
    snake_case    :      str      					=    (1 << p) - 1
    for _ in range(p - 2 ):
        snake_case    :      Any      					=    ((s * s) - 2) % m
    return s == 0


if __name__ == "__main__":
  print(lucas_lehmer_test(7))
  print(lucas_lehmer_test(11))





