
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     =  None
_A     =  {
    "7B": 1_10_08,
    "13B": 1_38_24,
    "30B": 1_79_20,
    "65B": 2_20_16,
    "70B": 2_86_72,
}
_A     =  {
    "7B": 1,
    "7Bf": 1,
    "13B": 2,
    "13Bf": 2,
    "30B": 4,
    "65B": 8,
    "70B": 8,
    "70Bf": 8,
}






def      lowerCamelCase__       (       __lowerCAmelCase						:						Optional[int]					,							__lowerCAmelCase						:						Union[str, Any]=1					,							__lowerCAmelCase						:						Dict=256       ):




			"""simple docstring"""


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






def      lowerCamelCase__       (       __lowerCAmelCase						:						List[Any]       ):




			"""simple docstring"""


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






def      lowerCamelCase__       (       __lowerCAmelCase						:						List[str]					,							__lowerCAmelCase						:						str       ):




			"""simple docstring"""


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






def      lowerCamelCase__       (       __lowerCAmelCase						:						str					,							__lowerCAmelCase						:						Optional[int]					,							__lowerCAmelCase						:						Union[str, Any]					,							__lowerCAmelCase						:						Optional[int]=True       ):




			"""simple docstring"""


			os.makedirs(__lowerCAmelCase					,							exist_ok=__lowerCAmelCase       )
			lowerCAmelCase_      			=			os.path.join(__lowerCAmelCase					,							"tmp"       )
			os.makedirs(__lowerCAmelCase					,							exist_ok=__lowerCAmelCase       )

			lowerCAmelCase_      			=			read_json(os.path.join(__lowerCAmelCase					,							"params.json"       )       )
			lowerCAmelCase_      			=			NUM_SHARDS[model_size]
			lowerCAmelCase_      			=			params["n_layers"]
			lowerCAmelCase_      			=			params["n_heads"]
			lowerCAmelCase_      			=			n_heads // num_shards
			lowerCAmelCase_      			=			params["dim"]
			lowerCAmelCase_      			=			dim // n_heads
			lowerCAmelCase_      			=			10_000.0
			lowerCAmelCase_      			=			1.0 / (base ** (torch.arange(0					,							__lowerCAmelCase					,							2       ).float() / dims_per_head))

			if "n_kv_heads" in params:
						lowerCAmelCase_      			=			params["n_kv_heads"]  # for GQA / MQA
						lowerCAmelCase_      			=			n_heads_per_shard // num_key_value_heads
						lowerCAmelCase_      			=			dim // num_key_value_heads
			else:  # compatibility with other checkpoints
						lowerCAmelCase_      			=			n_heads
						lowerCAmelCase_      			=			n_heads_per_shard
						lowerCAmelCase_      			=			dim

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

									lowerCAmelCase_      			=			{
									    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(),
									}
									lowerCAmelCase_      			=			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       )       )
									lowerCAmelCase_      			=			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					,							)
									lowerCAmelCase_      			=			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       )

									lowerCAmelCase_      			=			torch.cat(
									    [loaded[i][F"""layers.{layer_i}.attention.wo.weight"""] for i in range(__lowerCAmelCase       )]					,							dim=1       )
									lowerCAmelCase_      			=			torch.cat(
									    [loaded[i][F"""layers.{layer_i}.feed_forward.w1.weight"""] for i in range(__lowerCAmelCase       )]					,							dim=0       )
									lowerCAmelCase_      			=			torch.cat(
									    [loaded[i][F"""layers.{layer_i}.feed_forward.w2.weight"""] for i in range(__lowerCAmelCase       )]					,							dim=1       )
									lowerCAmelCase_      			=			torch.cat(
									    [loaded[i][F"""layers.{layer_i}.feed_forward.w3.weight"""] for i in range(__lowerCAmelCase       )]					,							dim=0       )

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

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

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




			"""simple docstring"""


			lowerCAmelCase_      			=			LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
			print(F"""Saving a {tokenizer_class.__name__} to {tokenizer_path}."""       )
			lowerCAmelCase_      			=			tokenizer_class(__lowerCAmelCase       )
			tokenizer.save_pretrained(__lowerCAmelCase       )






def      lowerCamelCase__       (       ):




			"""simple docstring"""


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


if __name__ == "__main__":
			main()





from ...configuration_utils import PretrainedConfig
from ...utils import logging


_A     =  logging.get_logger(__name__)

_A     =  {
    "transfo-xl-wt103": "https://huggingface.co/transfo-xl-wt103/resolve/main/config.json",
}





class 			_lowerCAmelCase      (			__a ):
				_lowercase    	='''transfo-xl'''
				_lowercase    	=['''mems''']
				_lowercase    	={
				    '''n_token''': '''vocab_size''',
				    '''hidden_size''': '''d_model''',
				    '''num_attention_heads''': '''n_head''',
				    '''num_hidden_layers''': '''n_layer''',
				}




				def __init__(      self       ,       _UpperCamelCase=267_735       ,       _UpperCamelCase=[20_000, 40_000, 200_000]       ,       _UpperCamelCase=1_024       ,       _UpperCamelCase=1_024       ,       _UpperCamelCase=16       ,       _UpperCamelCase=64       ,       _UpperCamelCase=4_096       ,       _UpperCamelCase=4       ,       _UpperCamelCase=False       ,       _UpperCamelCase=18       ,       _UpperCamelCase=1_600       ,       _UpperCamelCase=1_000       ,       _UpperCamelCase=True       ,       _UpperCamelCase=True       ,       _UpperCamelCase=0       ,       _UpperCamelCase=-1       ,       _UpperCamelCase=True       ,       _UpperCamelCase=0.1       ,       _UpperCamelCase=0.0       ,       _UpperCamelCase=True       ,       _UpperCamelCase="normal"       ,       _UpperCamelCase=0.01       ,       _UpperCamelCase=0.01       ,       _UpperCamelCase=0.02       ,       _UpperCamelCase=1e-5       ,       _UpperCamelCase=0       ,       **_UpperCamelCase       ,       )       ->   Dict:
							lowerCAmelCase_      			=			vocab_size
							lowerCAmelCase_      			=			[]
							self.cutoffs.extend(_UpperCamelCase  )
							if proj_share_all_but_first:
										lowerCAmelCase_      			=			[False] + [True] * len(self.cutoffs  )
							else:
										lowerCAmelCase_      			=			[False] + [False] * len(self.cutoffs  )
							lowerCAmelCase_      			=			d_model
							lowerCAmelCase_      			=			d_embed
							lowerCAmelCase_      			=			d_head
							lowerCAmelCase_      			=			d_inner
							lowerCAmelCase_      			=			div_val
							lowerCAmelCase_      			=			pre_lnorm
							lowerCAmelCase_      			=			n_layer
							lowerCAmelCase_      			=			n_head
							lowerCAmelCase_      			=			mem_len
							lowerCAmelCase_      			=			same_length
							lowerCAmelCase_      			=			attn_type
							lowerCAmelCase_      			=			clamp_len
							lowerCAmelCase_      			=			sample_softmax
							lowerCAmelCase_      			=			adaptive
							lowerCAmelCase_      			=			dropout
							lowerCAmelCase_      			=			dropatt
							lowerCAmelCase_      			=			untie_r
							lowerCAmelCase_      			=			init
							lowerCAmelCase_      			=			init_range
							lowerCAmelCase_      			=			proj_init_std
							lowerCAmelCase_      			=			init_std
							lowerCAmelCase_      			=			layer_norm_epsilon
							super().__init__(eos_token_id=_UpperCamelCase       ,       **_UpperCamelCase  )




				@property
				def   __a		(      self  )       ->   List[Any]:
							# Message copied from Transformer-XL documentation
							logger.info(f"""The model {self.model_type} is one of the few models that has no sequence length limit."""  )
							return -1




				@max_position_embeddings.setter
				def   __a		(      self       ,       _UpperCamelCase  )       ->   str:
							# Message copied from Transformer-XL documentation
							raise NotImplementedError(
							    f"""The model {self.model_type} is one of the few models that has no sequence length limit."""  )



