

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							:  Union[str, Any]            =						None
lowerCAmelCase							:  str            =						{
    """7B""": 11008,
    """13B""": 13824,
    """30B""": 17920,
    """65B""": 22016,
    """70B""": 28672,
}
lowerCAmelCase							:  Union[str, Any]            =						{
    """7B""": 1,
    """7Bf""": 1,
    """13B""": 2,
    """13Bf""": 2,
    """30B""": 4,
    """65B""": 8,
    """70B""": 8,
    """70Bf""": 8,
}







def       A_				(   _UpperCAmelCase    ,     _UpperCAmelCase=1    ,     _UpperCAmelCase=2_56   ):
						return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3   )   ) + multiple_of - 1) // multiple_of)







def       A_				(   _UpperCAmelCase   ):
						with open(_UpperCAmelCase    ,     "r"   ) as f:
												return json.load(_UpperCAmelCase   )







def       A_				(   _UpperCAmelCase    ,     _UpperCAmelCase   ):
						with open(_UpperCAmelCase    ,     "w"   ) as f:
												json.dump(_UpperCAmelCase    ,     _UpperCAmelCase   )







def       A_				(   _UpperCAmelCase    ,     _UpperCAmelCase    ,     _UpperCAmelCase    ,     _UpperCAmelCase=True   ):
						os.makedirs(_UpperCAmelCase    ,     exist_ok=_UpperCAmelCase   )
						SCREAMING_SNAKE_CASE_:			Optional[int]   		=	os.path.join(_UpperCAmelCase    ,     "tmp"   )
						os.makedirs(_UpperCAmelCase    ,     exist_ok=_UpperCAmelCase   )

						SCREAMING_SNAKE_CASE_:			Tuple   		=	read_json(os.path.join(_UpperCAmelCase    ,     "params.json"   )   )
						SCREAMING_SNAKE_CASE_:			Optional[Any]   		=	NUM_SHARDS[model_size]
						SCREAMING_SNAKE_CASE_:			int   		=	params["n_layers"]
						SCREAMING_SNAKE_CASE_:			Dict   		=	params["n_heads"]
						SCREAMING_SNAKE_CASE_:			Any   		=	n_heads // num_shards
						SCREAMING_SNAKE_CASE_:			Optional[Any]   		=	params["dim"]
						SCREAMING_SNAKE_CASE_:			Tuple   		=	dim // n_heads
						SCREAMING_SNAKE_CASE_:			Dict   		=	1_00_00.0
						SCREAMING_SNAKE_CASE_:			List[Any]   		=	1.0 / (base ** (torch.arange(0    ,     _UpperCAmelCase    ,     2   ).float() / dims_per_head))

						if "n_kv_heads" in params:
												SCREAMING_SNAKE_CASE_:			Dict   		=	params["n_kv_heads"]  # for GQA / MQA
												SCREAMING_SNAKE_CASE_:			Any   		=	n_heads_per_shard // num_key_value_heads
												SCREAMING_SNAKE_CASE_:			List[Any]   		=	dim // num_key_value_heads
						else:  # compatibility with other checkpoints
												SCREAMING_SNAKE_CASE_:			Tuple   		=	n_heads
												SCREAMING_SNAKE_CASE_:			Any   		=	n_heads_per_shard
												SCREAMING_SNAKE_CASE_:			Dict   		=	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.)
												SCREAMING_SNAKE_CASE_:			Tuple   		=	torch.load(os.path.join(_UpperCAmelCase    ,     "consolidated.00.pth"   )    ,     map_location="cpu"   )
						else:
												# Sharded
												SCREAMING_SNAKE_CASE_:			int   		=	[
												    torch.load(os.path.join(_UpperCAmelCase    ,     f"consolidated.{i:02d}.pth"   )    ,     map_location="cpu"   )
												    for i in range(_UpperCAmelCase   )
												]
						SCREAMING_SNAKE_CASE_:			Union[str, Any]   		=	0
						SCREAMING_SNAKE_CASE_:			int   		=	{"weight_map": {}}
						for layer_i in range(_UpperCAmelCase   ):
												SCREAMING_SNAKE_CASE_:			Optional[int]   		=	f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
												if model_size == "7B":
																		# Unsharded
																		SCREAMING_SNAKE_CASE_:			Optional[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.

																		SCREAMING_SNAKE_CASE_:			List[str]   		=	{
																		    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(),
																		}
																		SCREAMING_SNAKE_CASE_:			Any   		=	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   )   )
																		SCREAMING_SNAKE_CASE_:			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    ,     )
																		SCREAMING_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   )

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

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

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

						# Write configs
						SCREAMING_SNAKE_CASE_:			Any   		=	{"total_size": param_count * 2}
						write_json(_UpperCAmelCase    ,     os.path.join(_UpperCAmelCase    ,     "pytorch_model.bin.index.json"   )   )
						SCREAMING_SNAKE_CASE_:			Dict   		=	params["ffn_dim_multiplier"] if "ffn_dim_multiplier" in params else 1
						SCREAMING_SNAKE_CASE_:			List[Any]   		=	params["multiple_of"] if "multiple_of" in params else 2_56
						SCREAMING_SNAKE_CASE_:			int   		=	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."   )
						SCREAMING_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       A_				(   _UpperCAmelCase    ,     _UpperCAmelCase   ):
						# Initialize the tokenizer based on the `spm` model
						SCREAMING_SNAKE_CASE_:			int   		=	LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
						print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}."   )
						SCREAMING_SNAKE_CASE_:			Optional[Any]   		=	tokenizer_class(_UpperCAmelCase   )
						tokenizer.save_pretrained(_UpperCAmelCase   )







def       A_				(   ):
						SCREAMING_SNAKE_CASE_:			Optional[int]   		=	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`."   )
						SCREAMING_SNAKE_CASE_:			List[str]   		=	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    ,     )
						SCREAMING_SNAKE_CASE_:			Optional[int]   		=	os.path.join(args.input_dir    ,     "tokenizer.model"   )
						write_tokenizer(args.output_dir    ,     _UpperCAmelCase   )


if __name__ == "__main__":
						main()









import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_convbert import ConvBertTokenizer


lowerCAmelCase							:  Any            =						logging.get_logger(__name__)

lowerCAmelCase							:  Dict            =						{"""vocab_file""": """vocab.txt"""}

lowerCAmelCase							:  List[str]            =						{
    """vocab_file""": {
        """YituTech/conv-bert-base""": """https://huggingface.co/YituTech/conv-bert-base/resolve/main/vocab.txt""",
        """YituTech/conv-bert-medium-small""": (
            """https://huggingface.co/YituTech/conv-bert-medium-small/resolve/main/vocab.txt"""
        ),
        """YituTech/conv-bert-small""": """https://huggingface.co/YituTech/conv-bert-small/resolve/main/vocab.txt""",
    }
}

lowerCAmelCase							:  List[Any]            =						{
    """YituTech/conv-bert-base""": 512,
    """YituTech/conv-bert-medium-small""": 512,
    """YituTech/conv-bert-small""": 512,
}


lowerCAmelCase							:  Tuple            =						{
    """YituTech/conv-bert-base""": {"""do_lower_case""": True},
    """YituTech/conv-bert-medium-small""": {"""do_lower_case""": True},
    """YituTech/conv-bert-small""": {"""do_lower_case""": True},
}



class   __lowercase    (			UpperCAmelCase_     ):




			"""simple docstring"""







			_UpperCAmelCase       :							str      							=					VOCAB_FILES_NAMES
			_UpperCAmelCase       :							List[str]      							=					PRETRAINED_VOCAB_FILES_MAP
			_UpperCAmelCase       :							List[Any]      							=					PRETRAINED_INIT_CONFIGURATION
			_UpperCAmelCase       :							Optional[Any]      							=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			_UpperCAmelCase       :							Dict      							=					ConvBertTokenizer



			def __init__(			self  :					Optional[Any]       ,						lowerCAmelCase__  :					Optional[int]=None       ,						lowerCAmelCase__  :					List[str]=None       ,						lowerCAmelCase__  :					str=True       ,						lowerCAmelCase__  :					Any="[UNK]"       ,						lowerCAmelCase__  :					Optional[Any]="[SEP]"       ,						lowerCAmelCase__  :					Any="[PAD]"       ,						lowerCAmelCase__  :					Dict="[CLS]"       ,						lowerCAmelCase__  :					Dict="[MASK]"       ,						lowerCAmelCase__  :					Tuple=True       ,						lowerCAmelCase__  :					str=None       ,						**lowerCAmelCase__  :					Dict       ,						):
									super().__init__(
									    lowerCAmelCase__       ,						tokenizer_file=lowerCAmelCase__       ,						do_lower_case=lowerCAmelCase__       ,						unk_token=lowerCAmelCase__       ,						sep_token=lowerCAmelCase__       ,						pad_token=lowerCAmelCase__       ,						cls_token=lowerCAmelCase__       ,						mask_token=lowerCAmelCase__       ,						tokenize_chinese_chars=lowerCAmelCase__       ,						strip_accents=lowerCAmelCase__       ,						**lowerCAmelCase__       ,						)

									SCREAMING_SNAKE_CASE_:			List[str]  		=	json.loads(self.backend_tokenizer.normalizer.__getstate__())
									if (
									    normalizer_state.get("lowercase"       ,						lowerCAmelCase__) != do_lower_case
									    or normalizer_state.get("strip_accents"       ,						lowerCAmelCase__) != strip_accents
									    or normalizer_state.get("handle_chinese_chars"       ,						lowerCAmelCase__) != tokenize_chinese_chars
									):
															SCREAMING_SNAKE_CASE_:			Optional[int]  		=	getattr(lowerCAmelCase__       ,						normalizer_state.pop("type"))
															SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	do_lower_case
															SCREAMING_SNAKE_CASE_:			List[str]  		=	strip_accents
															SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	tokenize_chinese_chars
															SCREAMING_SNAKE_CASE_:			Optional[int]  		=	normalizer_class(**lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	do_lower_case



			def 	_SCREAMING_SNAKE_CASE      (			self  :					str       ,						lowerCAmelCase__  :					Optional[int]       ,						lowerCAmelCase__  :					Union[str, Any]=None):
									SCREAMING_SNAKE_CASE_:			List[str]  		=	[self.cls_token_id] + token_ids_a + [self.sep_token_id]

									if token_ids_a:
															output += token_ids_a + [self.sep_token_id]

									return output



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[Any]       ,						lowerCAmelCase__  :					List[int]       ,						lowerCAmelCase__  :					Optional[List[int]] = None):
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	[self.sep_token_id]
									SCREAMING_SNAKE_CASE_:			int  		=	[self.cls_token_id]
									if token_ids_a is None:
															return len(cls + token_ids_a + sep) * [0]
									return len(cls + token_ids_a + sep) * [0] + len(token_ids_a + sep) * [1]




			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					Optional[str] = None):
									SCREAMING_SNAKE_CASE_:			Any  		=	self._tokenizer.model.save(lowerCAmelCase__       ,						name=lowerCAmelCase__)
									return tuple(lowerCAmelCase__)






