


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"
		)
		__snake_case            =   None
__snake_case            =   {
    "7B": 11_008,
    "13B": 13_824,
    "30B": 17_920,
    "65B": 22_016,
    "70B": 28_672,
}
__snake_case            =   {
    "7B": 1,
    "7Bf": 1,
    "13B": 2,
    "13Bf": 2,
    "30B": 4,
    "65B": 8,
    "70B": 8,
    "70Bf": 8,
}






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		List[str]       ,  SCREAMING_SNAKE_CASE_   :		Dict=1       ,  SCREAMING_SNAKE_CASE_   :		Any=256	):





			"""simple docstring"""




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






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		List[Any]	):





			"""simple docstring"""




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






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		List[str]       ,  SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




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






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Tuple       ,  SCREAMING_SNAKE_CASE_   :		Union[str, Any]       ,  SCREAMING_SNAKE_CASE_   :		Tuple       ,  SCREAMING_SNAKE_CASE_   :		Optional[Any]=True	):





			"""simple docstring"""




			os.makedirs(SCREAMING_SNAKE_CASE_       ,  exist_ok=SCREAMING_SNAKE_CASE_	)
			UpperCamelCase         =					os.path.join(SCREAMING_SNAKE_CASE_       ,  """tmp"""	)
			os.makedirs(SCREAMING_SNAKE_CASE_       ,  exist_ok=SCREAMING_SNAKE_CASE_	)

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

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

			# permute for sliced rotary
			def permute(SCREAMING_SNAKE_CASE_   :		Optional[int]       ,  SCREAMING_SNAKE_CASE_   :		Tuple=n_heads       ,  SCREAMING_SNAKE_CASE_   :		List[str]=dim       ,  SCREAMING_SNAKE_CASE_   :		Union[str, Any]=dim	):
						return w.view(SCREAMING_SNAKE_CASE_       ,  dima // n_heads // 2       ,  2       ,  SCREAMING_SNAKE_CASE_	).transpose(1       ,  2	).reshape(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)

			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         =					torch.load(os.path.join(SCREAMING_SNAKE_CASE_       ,  """consolidated.00.pth"""	)       ,  map_location="""cpu"""	)
			else:
						# Sharded
						UpperCamelCase         =					[
						    torch.load(os.path.join(SCREAMING_SNAKE_CASE_       ,  f'consolidated.{i:02d}.pth'	)       ,  map_location="""cpu"""	)
						    for i in range(SCREAMING_SNAKE_CASE_	)
						]
			UpperCamelCase         =					0
			UpperCamelCase         =					{"""weight_map""": {}}
			for layer_i in range(SCREAMING_SNAKE_CASE_	):
						UpperCamelCase         =					f'pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin'
						if model_size == "7B":
									# Unsharded
									UpperCamelCase         =					{
									    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         =					{
									    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         =					permute(
									    torch.cat(
									        [
									            loaded[i][f'layers.{layer_i}.attention.wq.weight'].view(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)
									            for i in range(SCREAMING_SNAKE_CASE_	)
									        ]       ,  dim=0       ,  ).reshape(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)	)
									UpperCamelCase         =					permute(
									    torch.cat(
									        [
									            loaded[i][f'layers.{layer_i}.attention.wk.weight'].view(
									                SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)
									            for i in range(SCREAMING_SNAKE_CASE_	)
									        ]       ,  dim=0       ,  ).reshape(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  )
									UpperCamelCase         =					torch.cat(
									    [
									        loaded[i][f'layers.{layer_i}.attention.wv.weight'].view(
									            SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)
									        for i in range(SCREAMING_SNAKE_CASE_	)
									    ]       ,  dim=0       ,  ).reshape(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)

									UpperCamelCase         =					torch.cat(
									    [loaded[i][f'layers.{layer_i}.attention.wo.weight'] for i in range(SCREAMING_SNAKE_CASE_	)]       ,  dim=1	)
									UpperCamelCase         =					torch.cat(
									    [loaded[i][f'layers.{layer_i}.feed_forward.w1.weight'] for i in range(SCREAMING_SNAKE_CASE_	)]       ,  dim=0	)
									UpperCamelCase         =					torch.cat(
									    [loaded[i][f'layers.{layer_i}.feed_forward.w2.weight'] for i in range(SCREAMING_SNAKE_CASE_	)]       ,  dim=1	)
									UpperCamelCase         =					torch.cat(
									    [loaded[i][f'layers.{layer_i}.feed_forward.w3.weight'] for i in range(SCREAMING_SNAKE_CASE_	)]       ,  dim=0	)

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

			UpperCamelCase         =					f'pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin'
			if model_size == "7B":
						# Unsharded
						UpperCamelCase         =					{
						    """model.embed_tokens.weight""": loaded["""tok_embeddings.weight"""],
						    """model.norm.weight""": loaded["""norm.weight"""],
						    """lm_head.weight""": loaded["""output.weight"""],
						}
			else:
						UpperCamelCase         =					{
						    """model.norm.weight""": loaded[0]["""norm.weight"""],
						    """model.embed_tokens.weight""": torch.cat(
						        [loaded[i]["""tok_embeddings.weight"""] for i in range(SCREAMING_SNAKE_CASE_	)]       ,  dim=1	),
						    """lm_head.weight""": torch.cat([loaded[i]["""output.weight"""] for i in range(SCREAMING_SNAKE_CASE_	)]       ,  dim=0	),
						}

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

			# Write configs
			UpperCamelCase         =					{"""total_size""": param_count * 2}
			write_json(SCREAMING_SNAKE_CASE_       ,  os.path.join(SCREAMING_SNAKE_CASE_       ,  """pytorch_model.bin.index.json"""	)	)
			UpperCamelCase         =					params["""ffn_dim_multiplier"""] if """ffn_dim_multiplier""" in params else 1
			UpperCamelCase         =					params["""multiple_of"""] if """multiple_of""" in params else 256
			UpperCamelCase         =					LlamaConfig(
			    hidden_size=SCREAMING_SNAKE_CASE_       ,  intermediate_size=compute_intermediate_size(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	)       ,  num_attention_heads=params["""n_heads"""]       ,  num_hidden_layers=params["""n_layers"""]       ,  rms_norm_eps=params["""norm_eps"""]       ,  num_key_value_heads=SCREAMING_SNAKE_CASE_       ,  )
			config.save_pretrained(SCREAMING_SNAKE_CASE_	)

			# 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         =					LlamaForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE_       ,  torch_dtype=torch.floataa       ,  low_cpu_mem_usage=SCREAMING_SNAKE_CASE_	)
			# Avoid saving this as part of the config.
			del model.config._name_or_path

			print("""Saving in the Transformers format."""	)
			model.save_pretrained(SCREAMING_SNAKE_CASE_       ,  safe_serialization=SCREAMING_SNAKE_CASE_	)
			shutil.rmtree(SCREAMING_SNAKE_CASE_	)






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Optional[Any]       ,  SCREAMING_SNAKE_CASE_   :		Optional[Any]	):





			"""simple docstring"""




			UpperCamelCase         =					LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
			print(f'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'	)
			UpperCamelCase         =					tokenizer_class(SCREAMING_SNAKE_CASE_	)
			tokenizer.save_pretrained(SCREAMING_SNAKE_CASE_	)






def 							_lowercase    (		):





			"""simple docstring"""




			UpperCamelCase         =					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=SCREAMING_SNAKE_CASE_       ,  help="""Whether or not to save using `safetensors`."""	)
			UpperCamelCase         =					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         =					os.path.join(args.input_dir       ,  """tokenizer.model"""	)
			write_tokenizer(args.output_dir       ,  SCREAMING_SNAKE_CASE_	)


if __name__ == "__main__":
		main()








import warnings
from typing import List, Optional, Union

from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType




class    UpperCAmelCase					(       __snake_case    ):
			lowercase						= ["""image_processor""", """tokenizer"""]
			lowercase						= """FlavaImageProcessor"""
			lowercase						= ("""BertTokenizer""", """BertTokenizerFast""")




			def __init__(						self							:		Optional[int]       ,       __magic_name__							:		Optional[int]=None       ,       __magic_name__							:		str=None       ,       **__magic_name__							:		Any						):







						"""simple docstring"""







						UpperCamelCase         =					None
						if "feature_extractor" in kwargs:
									warnings.warn(
									    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
									    """ instead."""       ,       __magic_name__       ,       )
									UpperCamelCase         =					kwargs.pop("""feature_extractor"""						)

						UpperCamelCase         =					image_processor if image_processor is not None else feature_extractor
						if image_processor is None:
									raise ValueError("""You need to specify an `image_processor`."""						)
						if tokenizer is None:
									raise ValueError("""You need to specify a `tokenizer`."""						)

						super().__init__(__magic_name__       ,       __magic_name__						)
						UpperCamelCase         =					self.image_processor




			def __call__(						self							:		int       ,       __magic_name__							:		Optional[ImageInput] = None       ,       __magic_name__							:		Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None       ,       __magic_name__							:		bool = True       ,       __magic_name__							:		Union[bool, str, PaddingStrategy] = False       ,       __magic_name__							:		Union[bool, str, TruncationStrategy] = False       ,       __magic_name__							:		Optional[int] = None       ,       __magic_name__							:		int = 0       ,       __magic_name__							:		Optional[int] = None       ,       __magic_name__							:		Optional[bool] = None       ,       __magic_name__							:		Optional[bool] = None       ,       __magic_name__							:		Optional[bool] = None       ,       __magic_name__							:		Optional[bool] = None       ,       __magic_name__							:		bool = False       ,       __magic_name__							:		bool = False       ,       __magic_name__							:		bool = False       ,       __magic_name__							:		bool = False       ,       __magic_name__							:		bool = True       ,       __magic_name__							:		Optional[Union[str, TensorType]] = None       ,       **__magic_name__							:		str       ,       ):







						"""simple docstring"""







						if text is None and images is None:
									raise ValueError("""You have to specify either text or images. Both cannot be none."""						)

						if text is not None:
									UpperCamelCase         =					self.tokenizer(
									    text=__magic_name__       ,       add_special_tokens=__magic_name__       ,       padding=__magic_name__       ,       truncation=__magic_name__       ,       max_length=__magic_name__       ,       stride=__magic_name__       ,       pad_to_multiple_of=__magic_name__       ,       return_token_type_ids=__magic_name__       ,       return_attention_mask=__magic_name__       ,       return_overflowing_tokens=__magic_name__       ,       return_special_tokens_mask=__magic_name__       ,       return_offsets_mapping=__magic_name__       ,       return_length=__magic_name__       ,       verbose=__magic_name__       ,       return_tensors=__magic_name__       ,       **__magic_name__       ,       )
						if images is not None:
									UpperCamelCase         =					self.image_processor(
									    __magic_name__       ,       return_image_mask=__magic_name__       ,       return_codebook_pixels=__magic_name__       ,       return_tensors=__magic_name__       ,       **__magic_name__       ,       )

						if text is not None and images is not None:
									encoding.update(__magic_name__						)
									return encoding
						elif text is not None:
									return encoding
						else:
									return BatchEncoding(data=dict(**__magic_name__						)       ,       tensor_type=__magic_name__						)




			def   lowerCamelCase_					(						self							:		Tuple       ,       *__magic_name__							:		Dict       ,       **__magic_name__							:		List[Any]						):







						"""simple docstring"""







						return self.tokenizer.batch_decode(*__magic_name__       ,       **__magic_name__						)




			def   lowerCamelCase_					(						self							:		int       ,       *__magic_name__							:		Tuple       ,       **__magic_name__							:		Optional[Any]						):







						"""simple docstring"""







						return self.tokenizer.decode(*__magic_name__       ,       **__magic_name__						)




			@property
			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						UpperCamelCase         =					self.tokenizer.model_input_names
						UpperCamelCase         =					self.image_processor.model_input_names
						return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names						)						)




			@property
			def   lowerCamelCase_					(						self							:		int						):







						"""simple docstring"""







						warnings.warn(
						    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,       __magic_name__       ,       )
						return self.image_processor_class





			@property
			def   lowerCamelCase_					(						self							:		Tuple						):







						"""simple docstring"""







						warnings.warn(
						    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""       ,       __magic_name__       ,       )
						return self.image_processor




