


"""simple docstring"""






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       A_   (   _lowerCAmelCase     :					List[Any],     _lowerCAmelCase     :					Union[str, Any]=1,     _lowerCAmelCase     :					Any=2_56       ):






     """simple docstring"""

     return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3       )       ) + multiple_of - 1) // multiple_of)
def       A_   (   _lowerCAmelCase     :					int       ):






     """simple docstring"""

     with open(_lowerCAmelCase,     '''r'''       ) as f:
          return json.load(_lowerCAmelCase       )
def       A_   (   _lowerCAmelCase     :					Optional[Any],     _lowerCAmelCase     :					int       ):






     """simple docstring"""

     with open(_lowerCAmelCase,     '''w'''       ) as f:
          json.dump(_lowerCAmelCase,     _lowerCAmelCase       )
def       A_   (   _lowerCAmelCase     :					List[str],     _lowerCAmelCase     :					Optional[Any],     _lowerCAmelCase     :					Tuple,     _lowerCAmelCase     :					List[Any]=True       ):






     """simple docstring"""

     os.makedirs(_lowerCAmelCase,     exist_ok=_lowerCAmelCase       )
     _a    	=				os.path.join(_lowerCAmelCase,     '''tmp'''       )
     os.makedirs(_lowerCAmelCase,     exist_ok=_lowerCAmelCase       )

     _a    	=				read_json(os.path.join(_lowerCAmelCase,     '''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    	=				1_0_0_0_0.0
     _a    	=				1.0 / (base ** (torch.arange(0,     _lowerCAmelCase,     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(_lowerCAmelCase     :					List[str],     _lowerCAmelCase     :					Dict=n_heads,     _lowerCAmelCase     :					Optional[Any]=dim,     _lowerCAmelCase     :					int=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.)
          _a    	=				torch.load(os.path.join(_lowerCAmelCase,     '''consolidated.00.pth'''       ),     map_location='''cpu'''       )
     else:
          # Sharded
          _a    	=				[
              torch.load(os.path.join(_lowerCAmelCase,     f'consolidated.{i:02d}.pth'       ),     map_location='''cpu'''       )
              for i in range(_lowerCAmelCase       )
          ]
     _a    	=				0
     _a    	=				{'''weight_map''': {}}
     for layer_i in range(_lowerCAmelCase       ):
          _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(_lowerCAmelCase,     _lowerCAmelCase,     _lowerCAmelCase       )
                           for i in range(_lowerCAmelCase       )
                       ],     dim=0,     ).reshape(_lowerCAmelCase,     _lowerCAmelCase       )       )
               _a    	=				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,     )
               _a    	=				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       )

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

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

     _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(_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():
          _a    	=				filename
          param_count += v.numel()
     torch.save(_lowerCAmelCase,     os.path.join(_lowerCAmelCase,     _lowerCAmelCase       )       )

     # Write configs
     _a    	=				{'''total_size''': param_count * 2}
     write_json(_lowerCAmelCase,     os.path.join(_lowerCAmelCase,     '''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_56
     _a    	=				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.'''       )
     _a    	=				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     :					Union[str, Any],     _lowerCAmelCase     :					List[str]       ):






     """simple docstring"""

     _a    	=				LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
     print(f'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'       )
     _a    	=				tokenizer_class(_lowerCAmelCase       )
     tokenizer.save_pretrained(_lowerCAmelCase       )
def       A_   (   ):






     """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=_lowerCAmelCase,     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,     _lowerCAmelCase       )


if __name__ == "__main__":
   main()



"""simple docstring"""






def       A_   (   _lowerCAmelCase     :					list,     _lowerCAmelCase     :					list,     _lowerCAmelCase     :					int       ):






     """simple docstring"""

     _a    	=				len(_lowerCAmelCase       )
     _a    	=				[[0] * n for i in range(_lowerCAmelCase       )]
     for i in range(_lowerCAmelCase       ):
          _a    	=				y_points[i]

     for i in range(2,     _lowerCAmelCase       ):
          for j in range(_lowerCAmelCase,     _lowerCAmelCase       ):
               _a    	=				(
                   (xa - x_points[j - i + 1]) * q[j][i - 1]
                   - (xa - x_points[j]) * q[j - 1][i - 1]
               ) / (x_points[j] - x_points[j - i + 1])

     return [q[n - 1][n - 1], q]


if __name__ == "__main__":
   import doctest

   doctest.testmod()