
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_      			=			10000.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()





import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


_A     =  "▁"

_A     =  {"vocab_file": "spiece.model"}

_A     =  {
    "vocab_file": {"google/pegasus-xsum": "https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"}
}

_A     =  {
    "google/pegasus-xsum": 5_12,
}


_A     =  logging.get_logger(__name__)





class 			_lowerCAmelCase      (			__a ):
				_lowercase    	=VOCAB_FILES_NAMES

				_lowercase    	=VOCAB_FILES_NAMES
				_lowercase    	=PRETRAINED_VOCAB_FILES_MAP
				_lowercase    	=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				_lowercase    	=['''input_ids''', '''attention_mask''']




				def __init__(      self       ,       _UpperCamelCase       ,       _UpperCamelCase="<pad>"       ,       _UpperCamelCase="</s>"       ,       _UpperCamelCase="<unk>"       ,       _UpperCamelCase="<mask_2>"       ,       _UpperCamelCase="<mask_1>"       ,       _UpperCamelCase=None       ,       _UpperCamelCase=103       ,       _UpperCamelCase = None       ,       **_UpperCamelCase       ,       )       ->   None:
							lowerCAmelCase_      			=			offset
							if additional_special_tokens is not None:
										if not isinstance(_UpperCamelCase       ,       _UpperCamelCase  ):
													raise TypeError(
													    f"""additional_special_tokens should be of type {type(_UpperCamelCase  )}, but is"""
													    f""" {type(_UpperCamelCase  )}"""  )

										lowerCAmelCase_      			=			(
										    ([mask_token_sent] + additional_special_tokens)
										    if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
										    else additional_special_tokens
										)
										# fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
										additional_special_tokens_extended += [
										    f"""<unk_{i}>""" for i in range(len(_UpperCamelCase  )       ,       self.offset - 1  )
										]

										if len(set(_UpperCamelCase  )  ) != len(_UpperCamelCase  ):
													raise ValueError(
													    "Please make sure that the provided additional_special_tokens do not contain an incorrectly"
													    f""" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}."""  )
										lowerCAmelCase_      			=			additional_special_tokens_extended
							else:
										lowerCAmelCase_      			=			[mask_token_sent] if mask_token_sent is not None else []
										additional_special_tokens += [f"""<unk_{i}>""" for i in range(2       ,       self.offset  )]

							lowerCAmelCase_      			=			{} if sp_model_kwargs is None else sp_model_kwargs

							super().__init__(
							    eos_token=_UpperCamelCase       ,       unk_token=_UpperCamelCase       ,       mask_token=_UpperCamelCase       ,       pad_token=_UpperCamelCase       ,       mask_token_sent=_UpperCamelCase       ,       offset=_UpperCamelCase       ,       additional_special_tokens=_UpperCamelCase       ,       sp_model_kwargs=self.sp_model_kwargs       ,       **_UpperCamelCase       ,       )
							lowerCAmelCase_      			=			mask_token_sent
							lowerCAmelCase_      			=			vocab_file
							lowerCAmelCase_      			=			spm.SentencePieceProcessor(**self.sp_model_kwargs  )
							self.sp_model.Load(_UpperCamelCase  )

							# add special tokens to encoder dict
							lowerCAmelCase_      			=			{
							    0: self.pad_token,
							    1: self.eos_token,
							}

							if self.mask_token_sent is not None:
										self.encoder.update(
										    {
										        2: self.mask_token_sent,
										        3: self.mask_token,
										    }  )

							if self.offset > 0:
										# entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
										# mask_token_sent is already added to list -> so start at 1
										self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1       ,       self.offset - 1  )}  )

							lowerCAmelCase_      			=			{v: k for k, v in self.encoder.items()}




				@property
				def   __a		(      self  )       ->   int:
							return len(self.sp_model  ) + self.offset




				def   __a		(      self  )       ->   Dict[str, int]:
							lowerCAmelCase_      			=			{self.convert_ids_to_tokens(_UpperCamelCase  ): i for i in range(self.vocab_size  )}
							vocab.update(self.added_tokens_encoder  )
							return vocab




				def __getstate__(      self  )       ->   List[Any]:
							lowerCAmelCase_      			=			self.__dict__.copy()
							lowerCAmelCase_      			=			None
							return state




				def __setstate__(      self       ,       _UpperCamelCase  )       ->   Optional[int]:
							lowerCAmelCase_      			=			d

							# for backward compatibility
							if not hasattr(self       ,       "sp_model_kwargs"  ):
										lowerCAmelCase_      			=			{}

							lowerCAmelCase_      			=			spm.SentencePieceProcessor(**self.sp_model_kwargs  )
							self.sp_model.Load(self.vocab_file  )




				def   __a		(      self       ,       _UpperCamelCase  )       ->   List[str]:
							return self.sp_model.encode(_UpperCamelCase       ,       out_type=_UpperCamelCase  )




				def   __a		(      self       ,       _UpperCamelCase  )       ->   int:
							if token in self.decoder:
										return self.decoder[token]
							elif token in self.added_tokens_decoder:
										return self.added_tokens_decoder[token]
							lowerCAmelCase_      			=			self.sp_model.piece_to_id(_UpperCamelCase  )
							return sp_id + self.offset




				def   __a		(      self       ,       _UpperCamelCase  )       ->   str:
							if index in self.encoder:
										return self.encoder[index]
							elif index in self.added_tokens_encoder:
										return self.added_tokens_encoder[index]
							else:
										lowerCAmelCase_      			=			self.sp_model.IdToPiece(index - self.offset  )
							return token




				def   __a		(      self       ,       _UpperCamelCase  )       ->   Optional[Any]:
							lowerCAmelCase_      			=			[]
							lowerCAmelCase_      			=			""
							for token in tokens:
										# make sure that special tokens are not decoded using sentencepiece model
										if token in self.all_special_tokens:
													out_string += self.sp_model.decode(_UpperCamelCase  ) + token
													lowerCAmelCase_      			=			[]
										else:
													current_sub_tokens.append(_UpperCamelCase  )
							out_string += self.sp_model.decode(_UpperCamelCase  )
							return out_string.strip()




				def   __a		(      self       ,       _UpperCamelCase=False  )       ->   Optional[int]:
							return 1




				def   __a		(      self       ,       _UpperCamelCase  )       ->   int:
							lowerCAmelCase_      			=			set(self.all_special_ids  )  # call it once instead of inside list comp
							all_special_ids.remove(self.unk_token_id  )  # <unk> is only sometimes special

							return [1 if x in all_special_ids else 0 for x in seq]




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase = None       ,       _UpperCamelCase = False  )       ->   List[int]:
							if already_has_special_tokens:
										return self._special_token_mask(_UpperCamelCase  )
							elif token_ids_a is None:
										return self._special_token_mask(_UpperCamelCase  ) + [1]
							else:
										return self._special_token_mask(token_ids_a + token_ids_a  ) + [1]




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase=None  )       ->   List[int]:
							if token_ids_a is None:
										return token_ids_a + [self.eos_token_id]
							# We don't expect to process pairs, but leave the pair logic for API consistency
							return token_ids_a + token_ids_a + [self.eos_token_id]




				def   __a		(      self       ,       _UpperCamelCase       ,       _UpperCamelCase = None  )       ->   Tuple[str]:
							if not os.path.isdir(_UpperCamelCase  ):
										logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""  )
										return
							lowerCAmelCase_      			=			os.path.join(
							    _UpperCamelCase       ,       (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]  )

							if os.path.abspath(self.vocab_file  ) != os.path.abspath(_UpperCamelCase  ) and os.path.isfile(self.vocab_file  ):
										copyfile(self.vocab_file       ,       _UpperCamelCase  )
							elif not os.path.isfile(self.vocab_file  ):
										with open(_UpperCamelCase       ,       "wb"  ) as fi:
													lowerCAmelCase_      			=			self.sp_model.serialized_model_proto()
													fi.write(_UpperCamelCase  )

							return (out_vocab_file,)



