





import os
import string
import sys


lowerCAmelCase_    =    1 << 8

lowerCAmelCase_    =    {
    '''tab''': ord('''\t'''),
    '''newline''': ord('''\r'''),
    '''esc''': 2_7,
    '''up''': 6_5 + ARROW_KEY_FLAG,
    '''down''': 6_6 + ARROW_KEY_FLAG,
    '''right''': 6_7 + ARROW_KEY_FLAG,
    '''left''': 6_8 + ARROW_KEY_FLAG,
    '''mod_int''': 9_1,
    '''undefined''': sys.maxsize,
    '''interrupt''': 3,
    '''insert''': 5_0,
    '''delete''': 5_1,
    '''pg_up''': 5_3,
    '''pg_down''': 5_4,
}

lowerCAmelCase_    =    KEYMAP['''up''']
lowerCAmelCase_    =    KEYMAP['''left''']

if sys.platform == "win32":
	lowerCAmelCase_    =    []
	lowerCAmelCase_    =    {
	    B'''\xe0H''': KEYMAP['''up'''] - ARROW_KEY_FLAG,
	    B'''\x00H''': KEYMAP['''up'''] - ARROW_KEY_FLAG,
	    B'''\xe0P''': KEYMAP['''down'''] - ARROW_KEY_FLAG,
	    B'''\x00P''': KEYMAP['''down'''] - ARROW_KEY_FLAG,
	    B'''\xe0M''': KEYMAP['''right'''] - ARROW_KEY_FLAG,
	    B'''\x00M''': KEYMAP['''right'''] - ARROW_KEY_FLAG,
	    B'''\xe0K''': KEYMAP['''left'''] - ARROW_KEY_FLAG,
	    B'''\x00K''': KEYMAP['''left'''] - ARROW_KEY_FLAG,
	}

for i in range(1_0):
	lowerCAmelCase_    =    ord(str(i))






def 					lowerCamelCase_						(    )		->						Tuple:






				"""simple docstring"""


				if os.name == "nt":
								import msvcrt

								snake_case_						:   str           =  '''mbcs'''
								# Flush the keyboard buffer
								while msvcrt.kbhit():
												msvcrt.getch()
								if len(_UpperCamelCase			) == 0:
												# Read the keystroke
												snake_case_						:   Optional[int]           =  msvcrt.getch()

												# If it is a prefix char, get second part
												if ch in (b"\x00", b"\xe0"):
																snake_case_						:   List[str]           =  ch + msvcrt.getch()
																# Translate actual Win chars to bullet char types
																try:
																				snake_case_						:   Optional[int]           =  chr(WIN_KEYMAP[cha]			)
																				WIN_CH_BUFFER.append(chr(KEYMAP['''mod_int''']			)			)
																				WIN_CH_BUFFER.append(_UpperCamelCase			)
																				if ord(_UpperCamelCase			) in (
																				    KEYMAP["insert"] - 1 << 9,
																				    KEYMAP["delete"] - 1 << 9,
																				    KEYMAP["pg_up"] - 1 << 9,
																				    KEYMAP["pg_down"] - 1 << 9,
																				):
																								WIN_CH_BUFFER.append(chr(126			)			)
																				snake_case_						:   List[str]           =  chr(KEYMAP['''esc''']			)
																except KeyError:
																				snake_case_						:   Optional[Any]           =  cha[1]
												else:
																snake_case_						:   Optional[int]           =  ch.decode(_UpperCamelCase			)
								else:
												snake_case_						:   Optional[Any]           =  WIN_CH_BUFFER.pop(0			)
				elif os.name == "posix":
								import termios
								import tty

								snake_case_						:   Optional[Any]           =  sys.stdin.fileno()
								snake_case_						:   Dict           =  termios.tcgetattr(_UpperCamelCase			)
								try:
												tty.setraw(_UpperCamelCase			)
												snake_case_						:   List[Any]           =  sys.stdin.read(1			)
								finally:
												termios.tcsetattr(_UpperCamelCase			,    termios.TCSADRAIN			,    _UpperCamelCase			)
				return ch






def 					lowerCamelCase_						(    )		->						int:






				"""simple docstring"""


				snake_case_						:   List[Any]           =  get_raw_chars()
				if ord(_UpperCamelCase			) in [KEYMAP["interrupt"], KEYMAP["newline"]]:
								return char

				elif ord(_UpperCamelCase			) == KEYMAP["esc"]:
								snake_case_						:   Optional[Any]           =  get_raw_chars()
								if ord(_UpperCamelCase			) == KEYMAP["mod_int"]:
												snake_case_						:   str           =  get_raw_chars()
												if ord(_UpperCamelCase			) >= KEYMAP["arrow_begin"] - ARROW_KEY_FLAG and ord(_UpperCamelCase			) <= KEYMAP["arrow_end"] - ARROW_KEY_FLAG:
																return chr(ord(_UpperCamelCase			) + ARROW_KEY_FLAG			)
												else:
																return KEYMAP["undefined"]
								else:
												return get_raw_chars()

				else:
								if char in string.printable:
												return char
								else:
												return KEYMAP["undefined"]







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'''
	)
	lowerCAmelCase_    =    None
lowerCAmelCase_    =    {
    '''7B''': 1_1_0_0_8,
    '''13B''': 1_3_8_2_4,
    '''30B''': 1_7_9_2_0,
    '''65B''': 2_2_0_1_6,
    '''70B''': 2_8_6_7_2,
}
lowerCAmelCase_    =    {
    '''7B''': 1,
    '''7Bf''': 1,
    '''13B''': 2,
    '''13Bf''': 2,
    '''30B''': 4,
    '''65B''': 8,
    '''70B''': 8,
    '''70Bf''': 8,
}






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase=1			,    _UpperCamelCase=256			)		->						Optional[int]:






				"""simple docstring"""


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






def 					lowerCamelCase_						(    _UpperCamelCase			)		->						Optional[Any]:






				"""simple docstring"""


				with open(_UpperCamelCase			,    '''r'''			) as f:
								return json.load(_UpperCamelCase			)






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase			)		->						int:






				"""simple docstring"""


				with open(_UpperCamelCase			,    '''w'''			) as f:
								json.dump(_UpperCamelCase			,    _UpperCamelCase			)






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase			,    _UpperCamelCase			,    _UpperCamelCase=True			)		->						Optional[Any]:






				"""simple docstring"""


				os.makedirs(_UpperCamelCase			,    exist_ok=_UpperCamelCase			)
				snake_case_						:   int           =  os.path.join(_UpperCamelCase			,    '''tmp'''			)
				os.makedirs(_UpperCamelCase			,    exist_ok=_UpperCamelCase			)

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

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

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

												snake_case_						:   Union[str, 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(),
												}
												snake_case_						:   int           =  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			)			)
												snake_case_						:   Optional[int]           =  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			,    )
												snake_case_						:   int           =  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			)

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

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

				snake_case_						:   Any           =  f'''pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin'''
				if model_size == "7B":
								# Unsharded
								snake_case_						:   List[str]           =  {
								    '''model.embed_tokens.weight''': loaded['''tok_embeddings.weight'''],
								    '''model.norm.weight''': loaded['''norm.weight'''],
								    '''lm_head.weight''': loaded['''output.weight'''],
								}
				else:
								snake_case_						:   Dict           =  {
								    '''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():
								snake_case_						:   List[str]           =  filename
								param_count += v.numel()
				torch.save(_UpperCamelCase			,    os.path.join(_UpperCamelCase			,    _UpperCamelCase			)			)

				# Write configs
				snake_case_						:   int           =  {'''total_size''': param_count * 2}
				write_json(_UpperCamelCase			,    os.path.join(_UpperCamelCase			,    '''pytorch_model.bin.index.json'''			)			)
				snake_case_						:   str           =  params['''ffn_dim_multiplier'''] if '''ffn_dim_multiplier''' in params else 1
				snake_case_						:   Optional[int]           =  params['''multiple_of'''] if '''multiple_of''' in params else 256
				snake_case_						:   Optional[Any]           =  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.'''			)
				snake_case_						:   Union[str, 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			,    _UpperCamelCase			)		->						List[str]:






				"""simple docstring"""


				snake_case_						:   Union[str, Any]           =  LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
				print(f'''Saving a {tokenizer_class.__name__} to {tokenizer_path}.'''			)
				snake_case_						:   Union[str, Any]           =  tokenizer_class(_UpperCamelCase			)
				tokenizer.save_pretrained(_UpperCamelCase			)






def 					lowerCamelCase_						(    )		->						int:






				"""simple docstring"""


				snake_case_						:   Union[str, Any]           =  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`.'''			)
				snake_case_						:   Dict           =  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			,    )
				snake_case_						:   Dict           =  os.path.join(args.input_dir			,    '''tokenizer.model'''			)
				write_tokenizer(args.output_dir			,    _UpperCamelCase			)


if __name__ == "__main__":
	main()
