


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


def  _snake_case       (					A			,   A=1			,   A=256	)						->   List[Any]:
							return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3	)	) + multiple_of - 1) // multiple_of)


def  _snake_case       (					A	)						->   Optional[int]:
							with open(A			,   '''r'''	) as f:
														return json.load(A	)


def  _snake_case       (					A			,   A	)						->   Optional[Any]:
							with open(A			,   '''w'''	) as f:
														json.dump(A			,   A	)


def  _snake_case       (					A			,   A			,   A			,   A=True	)						->   Tuple:
							os.makedirs(A			,   exist_ok=A	)
							lowerCAmelCase__	       =     os.path.join(A			,   '''tmp'''	)
							os.makedirs(A			,   exist_ok=A	)

							lowerCAmelCase__	       =     read_json(os.path.join(A			,   '''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			,   A			,   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(A			,   A=n_heads			,   A=dim			,   A=dim	):
														return w.view(A			,   dima // n_heads // 2			,   2			,   A	).transpose(1			,   2	).reshape(A			,   A	)

							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(A			,   '''consolidated.00.pth'''	)			,   map_location='''cpu'''	)
							else:
														# Sharded
														lowerCAmelCase__	       =     [
														    torch.load(os.path.join(A			,   F"""consolidated.{i:02d}.pth"""	)			,   map_location='''cpu'''	)
														    for i in range(A	)
														]
							lowerCAmelCase__	       =     0
							lowerCAmelCase__	       =     {'''weight_map''': {}}
							for layer_i in range(A	):
														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(A			,   A			,   A	)
																					            for i in range(A	)
																					        ]			,   dim=0			,   ).reshape(A			,   A	)	)
																					lowerCAmelCase__	       =     permute(
																					    torch.cat(
																					        [
																					            loaded[i][F"""layers.{layer_i}.attention.wk.weight"""].view(
																					                A			,   A			,   A	)
																					            for i in range(A	)
																					        ]			,   dim=0			,   ).reshape(A			,   A	)			,   A			,   A			,   A			,   )
																					lowerCAmelCase__	       =     torch.cat(
																					    [
																					        loaded[i][F"""layers.{layer_i}.attention.wv.weight"""].view(
																					            A			,   A			,   A	)
																					        for i in range(A	)
																					    ]			,   dim=0			,   ).reshape(A			,   A	)

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

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

							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(A	)]			,   dim=1	),
														    '''lm_head.weight''': torch.cat([loaded[i]['''output.weight'''] for i in range(A	)]			,   dim=0	),
														}

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

							# Write configs
							lowerCAmelCase__	       =     {'''total_size''': param_count * 2}
							write_json(A			,   os.path.join(A			,   '''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=A			,   intermediate_size=compute_intermediate_size(A			,   A			,   A	)			,   num_attention_heads=params['''n_heads''']			,   num_hidden_layers=params['''n_layers''']			,   rms_norm_eps=params['''norm_eps''']			,   num_key_value_heads=A			,   )
							config.save_pretrained(A	)

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

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


def  _snake_case       (					A			,   A	)						->   Optional[Any]:
							# Initialize the tokenizer based on the `spm` model
							lowerCAmelCase__	       =     LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
							print(F"""Saving a {tokenizer_class.__name__} to {tokenizer_path}."""	)
							lowerCAmelCase__	       =     tokenizer_class(A	)
							tokenizer.save_pretrained(A	)


def  _snake_case       (					)						->   Union[str, Any]:
							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=A			,   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			,   A	)


if __name__ == "__main__":
		main()



'''simple docstring'''






from __future__ import annotations


def  _snake_case       (					A			,   A	)						->   float:
							lowerCAmelCase__	       =     sorted(numsa + numsa	)
							lowerCAmelCase__       , lowerCAmelCase__	       =     divmod(len(A	)			,   2	)
							if mod == 1:
														return all_numbers[div]
							else:
														return (all_numbers[div] + all_numbers[div - 1]) / 2


if __name__ == "__main__":
		import doctest

		doctest.testmod()
		__UpperCAmelCase       =							[float(x) for x in input('''Enter the elements of first array: ''').split()]
		__UpperCAmelCase       =							[float(x) for x in input('''Enter the elements of second array: ''').split()]
		print(f"""The median of two arrays is: {median_of_two_arrays(array_a, array_a)}""")