





'''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'
						)
						__lowerCamelCase						: List[str]		=   None
__lowerCamelCase						: Union[str, Any]		=   {
    '7B': 1_1008,
    '13B': 1_3824,
    '30B': 1_7920,
    '65B': 2_2016,
    '70B': 2_8672,
}
__lowerCamelCase						: Optional[Any]		=   {
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE=1  ,      __SCREAMING_SNAKE_CASE=256   ):


		"""simple docstring"""

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


def 					_a   (__SCREAMING_SNAKE_CASE   ):


		"""simple docstring"""

		with open(__SCREAMING_SNAKE_CASE  ,      '''r'''   ) as f:
				return json.load(__SCREAMING_SNAKE_CASE   )


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


		"""simple docstring"""

		with open(__SCREAMING_SNAKE_CASE  ,      '''w'''   ) as f:
				json.dump(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE=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       			=1_0000.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  ,      __SCREAMING_SNAKE_CASE=n_heads  ,      __SCREAMING_SNAKE_CASE=dim  ,      __SCREAMING_SNAKE_CASE=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 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


		"""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 					_a   ():


		"""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()













'''simple docstring'''






from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__lowerCamelCase						: Union[str, Any]		=   logging.get_logger(__name__)

__lowerCamelCase						: int		=   {
    'google/mobilenet_v2_1.4_224': 'https://huggingface.co/google/mobilenet_v2_1.4_224/resolve/main/config.json',
    'google/mobilenet_v2_1.0_224': 'https://huggingface.co/google/mobilenet_v2_1.0_224/resolve/main/config.json',
    'google/mobilenet_v2_0.75_160': 'https://huggingface.co/google/mobilenet_v2_0.75_160/resolve/main/config.json',
    'google/mobilenet_v2_0.35_96': 'https://huggingface.co/google/mobilenet_v2_0.35_96/resolve/main/config.json',
    # See all MobileNetV2 models at https://huggingface.co/models?filter=mobilenet_v2
}







class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				"""mobilenet_v2"""


  def __init__(				self    :					Tuple							,	UpperCamelCase__    :					Union[str, Any]=3							,	UpperCamelCase__    :					Dict=224							,	UpperCamelCase__    :					str=1.0							,	UpperCamelCase__    :					List[Any]=8							,	UpperCamelCase__    :					Union[str, Any]=8							,	UpperCamelCase__    :					str=6							,	UpperCamelCase__    :					str=32							,	UpperCamelCase__    :					str=True							,	UpperCamelCase__    :					int=True							,	UpperCamelCase__    :					str="relu6"							,	UpperCamelCase__    :					Tuple=True							,	UpperCamelCase__    :					List[str]=0.8							,	UpperCamelCase__    :					List[str]=0.02							,	UpperCamelCase__    :					str=0.001							,	UpperCamelCase__    :					Dict=255							,	**UpperCamelCase__    :					Tuple							,	)     ->						List[Any]:
    super().__init__(**UpperCamelCase__  )

    if depth_multiplier <= 0:
      raise ValueError('''depth_multiplier must be greater than zero.'''  )

    _UpperCamelCase       			=num_channels
    _UpperCamelCase       			=image_size
    _UpperCamelCase       			=depth_multiplier
    _UpperCamelCase       			=depth_divisible_by
    _UpperCamelCase       			=min_depth
    _UpperCamelCase       			=expand_ratio
    _UpperCamelCase       			=output_stride
    _UpperCamelCase       			=first_layer_is_expansion
    _UpperCamelCase       			=finegrained_output
    _UpperCamelCase       			=hidden_act
    _UpperCamelCase       			=tf_padding
    _UpperCamelCase       			=classifier_dropout_prob
    _UpperCamelCase       			=initializer_range
    _UpperCamelCase       			=layer_norm_eps
    _UpperCamelCase       			=semantic_loss_ignore_index







class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				version.parse("""1.11""")


  @property
  def   UpperCamelCase__      (				self    :					str  )     ->						Mapping[str, Mapping[int, str]]:
    return OrderedDict([('''pixel_values''', {0: '''batch'''})]  )


  @property
  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Mapping[str, Mapping[int, str]]:
    if self.task == "image-classification":
      return OrderedDict([('''logits''', {0: '''batch'''})]  )
    else:
      return OrderedDict([('''last_hidden_state''', {0: '''batch'''}), ('''pooler_output''', {0: '''batch'''})]  )




  @property
  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						float:
    return 1E-4






