





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 	lowerCAmelCase_					(						__UpperCAmelCase:    List[str]	,     __UpperCAmelCase:    Dict=1	,     __UpperCAmelCase:    Optional[Any]=256   )   ->			int:
    return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3   )   ) + multiple_of - 1) // multiple_of)




def 	lowerCAmelCase_					(						__UpperCAmelCase:    int   )   ->			Optional[int]:
    with open(__UpperCAmelCase	,     '''r'''   ) as f:
        return json.load(__UpperCAmelCase   )




def 	lowerCAmelCase_					(						__UpperCAmelCase:    int	,     __UpperCAmelCase:    List[Any]   )   ->			Dict:
    with open(__UpperCAmelCase	,     '''w'''   ) as f:
        json.dump(__UpperCAmelCase	,     __UpperCAmelCase   )




def 	lowerCAmelCase_					(						__UpperCAmelCase:    str	,     __UpperCAmelCase:    List[Any]	,     __UpperCAmelCase:    Dict	,     __UpperCAmelCase:    int=True   )   ->			Tuple:
    os.makedirs(__UpperCAmelCase	,     exist_ok=__UpperCAmelCase   )
    UpperCamelCase__							:       Union[str, Any]      =   os.path.join(__UpperCAmelCase	,     '''tmp'''   )
    os.makedirs(__UpperCAmelCase	,     exist_ok=__UpperCAmelCase   )

    UpperCamelCase__							:       Any      =   read_json(os.path.join(__UpperCAmelCase	,     '''params.json'''   )   )
    UpperCamelCase__							:       Optional[int]      =   NUM_SHARDS[model_size]
    UpperCamelCase__							:       str      =   params['''n_layers''']
    UpperCamelCase__							:       Any      =   params['''n_heads''']
    UpperCamelCase__							:       List[Any]      =   n_heads // num_shards
    UpperCamelCase__							:       Any      =   params['''dim''']
    UpperCamelCase__							:       int      =   dim // n_heads
    UpperCamelCase__							:       Any      =   10000.0
    UpperCamelCase__							:       int      =   1.0 / (base ** (torch.arange(0	,     __UpperCAmelCase	,     2   ).float() / dims_per_head))

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

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

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

            UpperCamelCase__							:       int      =   {
                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(),
            }
            UpperCamelCase__							:       Optional[Any]      =   permute(
                torch.cat(
                    [
                        loaded[i][f"layers.{layer_i}.attention.wq.weight"].view(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )
                        for i in range(__UpperCAmelCase   )
                    ]	,     dim=0	,     ).reshape(__UpperCAmelCase	,     __UpperCAmelCase   )   )
            UpperCamelCase__							:       str      =   permute(
                torch.cat(
                    [
                        loaded[i][f"layers.{layer_i}.attention.wk.weight"].view(
                            __UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )
                        for i in range(__UpperCAmelCase   )
                    ]	,     dim=0	,     ).reshape(__UpperCAmelCase	,     __UpperCAmelCase   )	,     __UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase	,     )
            UpperCamelCase__							:       Union[str, Any]      =   torch.cat(
                [
                    loaded[i][f"layers.{layer_i}.attention.wv.weight"].view(
                        __UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   )
                    for i in range(__UpperCAmelCase   )
                ]	,     dim=0	,     ).reshape(__UpperCAmelCase	,     __UpperCAmelCase   )

            UpperCamelCase__							:       Any      =   torch.cat(
                [loaded[i][f"layers.{layer_i}.attention.wo.weight"] for i in range(__UpperCAmelCase   )]	,     dim=1   )
            UpperCamelCase__							:       Tuple      =   torch.cat(
                [loaded[i][f"layers.{layer_i}.feed_forward.w1.weight"] for i in range(__UpperCAmelCase   )]	,     dim=0   )
            UpperCamelCase__							:       List[Any]      =   torch.cat(
                [loaded[i][f"layers.{layer_i}.feed_forward.w2.weight"] for i in range(__UpperCAmelCase   )]	,     dim=1   )
            UpperCamelCase__							:       Optional[Any]      =   torch.cat(
                [loaded[i][f"layers.{layer_i}.feed_forward.w3.weight"] for i in range(__UpperCAmelCase   )]	,     dim=0   )

        UpperCamelCase__							:       Optional[Any]      =   inv_freq
        for k, v in state_dict.items():
            UpperCamelCase__							:       List[Any]      =   filename
            param_count += v.numel()
        torch.save(__UpperCAmelCase	,     os.path.join(__UpperCAmelCase	,     __UpperCAmelCase   )   )

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

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

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

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

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




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Dict	,     __UpperCAmelCase:    Optional[int]   )   ->			Tuple:
    # Initialize the tokenizer based on the `spm` model
    UpperCamelCase__							:       List[str]      =   LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
    print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}."   )
    UpperCamelCase__							:       List[Any]      =   tokenizer_class(__UpperCAmelCase   )
    tokenizer.save_pretrained(__UpperCAmelCase   )




def 	lowerCAmelCase_					(						)   ->			Dict:
    UpperCamelCase__							:       Optional[int]      =   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=__UpperCAmelCase	,     help='''Whether or not to save using `safetensors`.'''   )
    UpperCamelCase__							:       List[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	,     )
    UpperCamelCase__							:       Tuple      =   os.path.join(args.input_dir	,     '''tokenizer.model'''   )
    write_tokenizer(args.output_dir	,     __UpperCAmelCase   )


if __name__ == "__main__":
  main()








import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Union

from .generation.configuration_utils import GenerationConfig
from .training_args import TrainingArguments
from .utils import add_start_docstrings


UpperCAmelCase_  							=				logging.getLogger(__name__)






@dataclass
@add_start_docstrings(TrainingArguments.__doc__       )
class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							bool     	=	field(default=__lowerCamelCase    ,       metadata={"help": "Whether to use SortishSampler or not."}       )
   a		:							bool     	=	field(
       default=__lowerCamelCase    ,       metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}       )
   a		:							Optional[int]     	=	field(
       default=__lowerCamelCase    ,       metadata={
           "help": (
               "The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default "
               "to the `max_length` value of the model configuration."
           )
       }    ,       )
   a		:							Optional[int]     	=	field(
       default=__lowerCamelCase    ,       metadata={
           "help": (
               "The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default "
               "to the `num_beams` value of the model configuration."
           )
       }    ,       )
   a		:							Optional[Union[str, Path, GenerationConfig]]     	=	field(
       default=__lowerCamelCase    ,       metadata={
           "help": "Model id, file path or url pointing to a GenerationConfig json file, to use during prediction."
       }    ,       )




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       Tuple      =   super().to_dict()
       for k, v in d.items():
           if isinstance(__magic_name__,	__magic_name__							):
               UpperCamelCase__							:       Tuple      =   v.to_dict()
       return d

