


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         =					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_   :		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 math
import unittest






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			assert isinstance(SCREAMING_SNAKE_CASE_       ,  SCREAMING_SNAKE_CASE_	) and (
			    number >= 0
			), "'number' must been an int and positive"

			if 1 < number < 4:
						# 2 and 3 are primes
						return True
			elif number < 2 or number % 2 == 0 or number % 3 == 0:
						# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
						return False

			# All primes number are in format of 6k +/- 1
			for i in range(5       ,  int(math.sqrt(SCREAMING_SNAKE_CASE_	) + 1	)       ,  6	):
						if number % i == 0 or number % (i + 2) == 0:
									return False
			return True




class    UpperCAmelCase					(       unittest.TestCase    ):




			def   lowerCamelCase_					(						self							:		Any						):







						"""simple docstring"""







						self.assertTrue(is_prime(2						)						)
						self.assertTrue(is_prime(3						)						)
						self.assertTrue(is_prime(5						)						)
						self.assertTrue(is_prime(7						)						)
						self.assertTrue(is_prime(1_1						)						)
						self.assertTrue(is_prime(1_3						)						)
						self.assertTrue(is_prime(1_7						)						)
						self.assertTrue(is_prime(1_9						)						)
						self.assertTrue(is_prime(2_3						)						)
						self.assertTrue(is_prime(2_9						)						)





			def   lowerCamelCase_					(						self							:		Dict						):







						"""simple docstring"""







						with self.assertRaises(__magic_name__						):
									is_prime(-1_9						)
						self.assertFalse(
						    is_prime(0						)       ,       """Zero doesn't have any positive factors, primes must have exactly two."""       ,       )
						self.assertFalse(
						    is_prime(1						)       ,       """One only has 1 positive factor, primes must have exactly two."""       ,       )
						self.assertFalse(is_prime(2 * 2						)						)
						self.assertFalse(is_prime(2 * 3						)						)
						self.assertFalse(is_prime(3 * 3						)						)
						self.assertFalse(is_prime(3 * 5						)						)
						self.assertFalse(is_prime(3 * 5 * 7						)						)


if __name__ == "__main__":
			unittest.main()




