






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"
  )
  __a      :       List[Any]				  =   None
__a      :       List[str]				  =   {
    "7B": 11008,
    "13B": 13824,
    "30B": 17920,
    "65B": 22016,
    "70B": 28672,
}
__a      :       Union[str, Any]				  =   {
    "7B": 1,
    "7Bf": 1,
    "13B": 2,
    "13Bf": 2,
    "30B": 4,
    "65B": 8,
    "70B": 8,
    "70Bf": 8,
}



def 			lowerCamelCase__						( __lowercase   :			Optional[int]			,				__lowercase   :			Tuple=1			,				__lowercase   :			Union[str, Any]=2_5_6    )						->      int:





    """simple docstring"""

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



def 			lowerCamelCase__						( __lowercase   :			int    )						->      List[str]:





    """simple docstring"""

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



def 			lowerCamelCase__						( __lowercase   :			Union[str, Any]			,				__lowercase   :			Any    )						->      str:





    """simple docstring"""

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



def 			lowerCamelCase__						( __lowercase   :			Union[str, Any]			,				__lowercase   :			Union[str, Any]			,				__lowercase   :			Union[str, Any]			,				__lowercase   :			Optional[int]=True    )						->      List[str]:





    """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         =	10_000.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   :			Tuple			,				__lowercase   :			List[str]=n_heads			,				__lowercase   :			Dict=dim			,				__lowercase   :			Union[str, Any]=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 2_5_6
    __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 			lowerCamelCase__						( __lowercase   :			Dict			,				__lowercase   :			Dict    )						->      Dict:





    """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 			lowerCamelCase__						( )						->      int:





    """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 typing import TYPE_CHECKING

from ...utils import _LazyModule


__a      :       Optional[int]				  =   {"tokenization_bertweet": ["BertweetTokenizer"]}


if TYPE_CHECKING:
  from .tokenization_bertweet import BertweetTokenizer

else:
  import sys

  __a      :       int				  =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




