



'''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'''
					)
					_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      lowercase_		(      lowerCAmelCase__		:     int	,							lowerCAmelCase__		:     Optional[Any]=1	,							lowerCAmelCase__		:     List[str]=256       ):

			"""simple docstring"""





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


def      lowercase_		(      lowerCAmelCase__		:     str       ):

			"""simple docstring"""





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


def      lowercase_		(      lowerCAmelCase__		:     Any	,							lowerCAmelCase__		:     Tuple       ):

			"""simple docstring"""





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


def      lowercase_		(      lowerCAmelCase__		:     List[str]	,							lowerCAmelCase__		:     Any	,							lowerCAmelCase__		:     Tuple	,							lowerCAmelCase__		:     str=True       ):

			"""simple docstring"""





			os.makedirs(__lowerCAmelCase	,							exist_ok=__lowerCAmelCase       )
			__UpperCAmelCase      :					Union[str, Any]   =   os.path.join(__lowerCAmelCase	,							"""tmp"""       )
			os.makedirs(__lowerCAmelCase	,							exist_ok=__lowerCAmelCase       )

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

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

			# permute for sliced rotary
			def permute(lowerCAmelCase__		:     str	,							lowerCAmelCase__		:     Tuple=n_heads	,							lowerCAmelCase__		:     Dict=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.)
						__UpperCAmelCase      :					Dict   =   torch.load(os.path.join(__lowerCAmelCase	,							"""consolidated.00.pth"""       )	,							map_location="""cpu"""       )
			else:
						# Sharded
						__UpperCAmelCase      :					Optional[int]   =   [
						    torch.load(os.path.join(__lowerCAmelCase	,							f'consolidated.{i:02d}.pth'       )	,							map_location="""cpu"""       )
						    for i in range(__lowerCAmelCase       )
						]
			__UpperCAmelCase      :					Optional[Any]   =   0
			__UpperCAmelCase      :					str   =   {"weight_map": {}}
			for layer_i in range(__lowerCAmelCase       ):
						__UpperCAmelCase      :					str   =   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      :					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(),
									}
									__UpperCAmelCase      :					Optional[int]   =   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       )       )
									__UpperCAmelCase      :					Optional[Any]   =   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	,							)
									__UpperCAmelCase      :					Dict   =   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       )

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

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

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

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

			"""simple docstring"""





			__UpperCAmelCase      :					List[Any]   =   LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
			print(f'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'       )
			__UpperCAmelCase      :					Dict   =   tokenizer_class(__lowerCAmelCase       )
			tokenizer.save_pretrained(__lowerCAmelCase       )


def      lowercase_		(      ):

			"""simple docstring"""





			__UpperCAmelCase      :					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=__lowerCAmelCase	,							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      :					int   =   os.path.join(args.input_dir	,							"""tokenizer.model"""       )
			write_tokenizer(args.output_dir	,							__lowerCAmelCase       )


if __name__ == "__main__":
					main()








'''simple docstring'''


import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):
		assert isinstance(__lowerCAmelCase					,       __lowerCAmelCase      )
		assert dataset.num_rows == 4
		assert dataset.num_columns == 1
		assert dataset.column_names == ["text"]
		for feature, expected_dtype in expected_features.items():
				assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("keep_in_memory"					,       [False, True]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Union[str, Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				int	    =					{"text": "string"}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase					,       keep_in_memory=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize(
    "features"					,       [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]					,       )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Any	    =					{"text": "string"}
		_UpperCAmelCase						:				Optional[Any]	    =					features.copy() if features else default_expected_features
		_UpperCAmelCase						:				Union[str, Any]	    =					(
		    Features({feature: Value(__lowerCAmelCase      ) for feature, dtype in features.items()}      ) if features is not None else None
		)
		_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       features=__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize("split"					,       [None, NamedSplit("train"      ), "train", "test"]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Any	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Dict	    =					{"text": "string"}
		_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase					,       split=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )
		assert dataset.split == split if split else "train"






@pytest.mark.parametrize("path_type"					,       [str, list]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		if issubclass(__lowerCAmelCase					,       __lowerCAmelCase      ):
				_UpperCAmelCase						:				Union[str, Any]	    =					text_path
		elif issubclass(__lowerCAmelCase					,       __lowerCAmelCase      ):
				_UpperCAmelCase						:				int	    =					[text_path]
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Union[str, Any]	    =					{"text": "string"}
		_UpperCAmelCase						:				Optional[int]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase=("train",)      ):
		assert isinstance(__lowerCAmelCase					,       __lowerCAmelCase      )
		for split in splits:
				_UpperCAmelCase						:				List[str]	    =					dataset_dict[split]
				assert dataset.num_rows == 4
				assert dataset.num_columns == 1
				assert dataset.column_names == ["text"]
				for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("keep_in_memory"					,       [False, True]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Tuple	    =					{"text": "string"}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCAmelCase						:				Any	    =					TextDatasetReader({"train": text_path}					,       cache_dir=__lowerCAmelCase					,       keep_in_memory=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize(
    "features"					,       [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]					,       )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
		_UpperCAmelCase						:				List[Any]	    =					{"text": "string"}
		_UpperCAmelCase						:				List[str]	    =					features.copy() if features else default_expected_features
		_UpperCAmelCase						:				Optional[int]	    =					(
		    Features({feature: Value(__lowerCAmelCase      ) for feature, dtype in features.items()}      ) if features is not None else None
		)
		_UpperCAmelCase						:				Tuple	    =					TextDatasetReader({"train": text_path}					,       features=__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize("split"					,       [None, NamedSplit("train"      ), "train", "test"]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		if split:
				_UpperCAmelCase						:				int	    =					{split: text_path}
		else:
				_UpperCAmelCase						:				Tuple	    =					"train"
				_UpperCAmelCase						:				List[str]	    =					{"train": text_path, "test": text_path}
		_UpperCAmelCase						:				Optional[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Optional[int]	    =					{"text": "string"}
		_UpperCAmelCase						:				int	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase					,       splits=list(path.keys()      )      )
		assert all(dataset[split].split == split for split in path.keys()      )



