


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_   :		Optional[Any]       ,  SCREAMING_SNAKE_CASE_   :		int=1       ,  SCREAMING_SNAKE_CASE_   :		Dict=256	):





   """simple docstring"""




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






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		Any	):





   """simple docstring"""




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






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





   """simple docstring"""




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






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





   """simple docstring"""




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

   UpperCamelCase         =					read_json(os.path.join(__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       ,  __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_   :		Union[str, Any]       ,  SCREAMING_SNAKE_CASE_   :		Dict=n_heads       ,  SCREAMING_SNAKE_CASE_   :		Optional[Any]=dim       ,  SCREAMING_SNAKE_CASE_   :		int=dim	):
      return w.view(__snake_case       ,  dima // n_heads // 2       ,  2       ,  __snake_case	).transpose(1       ,  2	).reshape(__snake_case       ,  __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(__snake_case       ,  """consolidated.00.pth"""	)       ,  map_location="""cpu"""	)
   else:
      # Sharded
      UpperCamelCase         =					[
          torch.load(os.path.join(__snake_case       ,  f'consolidated.{i:02d}.pth'	)       ,  map_location="""cpu"""	)
          for i in range(__snake_case	)
      ]
   UpperCamelCase         =					0
   UpperCamelCase         =					{"""weight_map""": {}}
   for layer_i in range(__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(__snake_case       ,  __snake_case       ,  __snake_case	)
                     for i in range(__snake_case	)
                 ]       ,  dim=0       ,  ).reshape(__snake_case       ,  __snake_case	)	)
         UpperCamelCase         =					permute(
             torch.cat(
                 [
                     loaded[i][f'layers.{layer_i}.attention.wk.weight'].view(
                         __snake_case       ,  __snake_case       ,  __snake_case	)
                     for i in range(__snake_case	)
                 ]       ,  dim=0       ,  ).reshape(__snake_case       ,  __snake_case	)       ,  __snake_case       ,  __snake_case       ,  __snake_case       ,  )
         UpperCamelCase         =					torch.cat(
             [
                 loaded[i][f'layers.{layer_i}.attention.wv.weight'].view(
                     __snake_case       ,  __snake_case       ,  __snake_case	)
                 for i in range(__snake_case	)
             ]       ,  dim=0       ,  ).reshape(__snake_case       ,  __snake_case	)

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

      UpperCamelCase         =					inv_freq
      for k, v in state_dict.items():
         UpperCamelCase         =					filename
         param_count += v.numel()
      torch.save(__snake_case       ,  os.path.join(__snake_case       ,  __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(__snake_case	)]       ,  dim=1	),
          """lm_head.weight""": torch.cat([loaded[i]["""output.weight"""] for i in range(__snake_case	)]       ,  dim=0	),
      }

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

   # Write configs
   UpperCamelCase         =					{"""total_size""": param_count * 2}
   write_json(__snake_case       ,  os.path.join(__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=__snake_case       ,  intermediate_size=compute_intermediate_size(__snake_case       ,  __snake_case       ,  __snake_case	)       ,  num_attention_heads=params["""n_heads"""]       ,  num_hidden_layers=params["""n_layers"""]       ,  rms_norm_eps=params["""norm_eps"""]       ,  num_key_value_heads=__snake_case       ,  )
   config.save_pretrained(__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(__snake_case       ,  torch_dtype=torch.floataa       ,  low_cpu_mem_usage=__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(__snake_case       ,  safe_serialization=__snake_case	)
   shutil.rmtree(__snake_case	)






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





   """simple docstring"""




   UpperCamelCase         =					LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
   print(f'Saving a {tokenizer_class.__name__} to {tokenizer_path}.'	)
   UpperCamelCase         =					tokenizer_class(__snake_case	)
   tokenizer.save_pretrained(__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=__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       ,  __snake_case	)


if __name__ == "__main__":
   main()








import math






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int	):





			"""simple docstring"""




			UpperCamelCase         =					[True] * n
			UpperCamelCase         =					False
			UpperCamelCase         =					False
			UpperCamelCase         =					True

			for i in range(3       ,  int(n**0.5 + 1	)       ,  2	):
						UpperCamelCase         =					i * 2
						while index < n:
									UpperCamelCase         =					False
									UpperCamelCase         =					index + i

			UpperCamelCase         =					[2]

			for i in range(3       ,  SCREAMING_SNAKE_CASE_       ,  2	):
						if is_prime[i]:
									primes.append(SCREAMING_SNAKE_CASE_	)

			return primes






def 							_lowercase    (		SCREAMING_SNAKE_CASE_   :		int = 999_966_663_333	):





			"""simple docstring"""




			UpperCamelCase         =					math.floor(math.sqrt(SCREAMING_SNAKE_CASE_	)	) + 100
			UpperCamelCase         =					prime_sieve(SCREAMING_SNAKE_CASE_	)

			UpperCamelCase         =					0
			UpperCamelCase         =					0
			UpperCamelCase         =					primes[prime_index]

			while (last_prime**2) <= limit:
						UpperCamelCase         =					primes[prime_index + 1]

						UpperCamelCase         =					last_prime**2
						UpperCamelCase         =					next_prime**2

						# Get numbers divisible by lps(current)
						UpperCamelCase         =					lower_bound + last_prime
						while upper_bound > current <= limit:
									matches_sum += current
									current += last_prime

						# Reset the upper_bound
						while (upper_bound - next_prime) > limit:
									upper_bound -= next_prime

						# Add the numbers divisible by ups(current)
						UpperCamelCase         =					upper_bound - next_prime
						while current > lower_bound:
									matches_sum += current
									current -= next_prime

						# Remove the numbers divisible by both ups and lps
						UpperCamelCase         =					0
						while upper_bound > current <= limit:
									if current <= lower_bound:
												# Increment the current number
												current += last_prime * next_prime
												continue

									if current > limit:
												break

									# Remove twice since it was added by both ups and lps
									matches_sum -= current * 2

									# Increment the current number
									current += last_prime * next_prime

						# Setup for next pair
						UpperCamelCase         =					next_prime
						prime_index += 1

			return matches_sum


if __name__ == "__main__":
			print(solution())




