
import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer
class 						_A :






  def __init__(  self		:     List[str]				,						_A		:     str				,						_A		:     List[str]=13				,						_A		:     List[str]=7				,						_A		:     Dict=True				,						_A		:     Any=True				,						_A		:     int=False				,						_A		:     str=True				,						_A		:     str=99				,						_A		:     Dict=32				,						_A		:     str=5				,						_A		:     Any=4				,						_A		:     Dict=37				,						_A		:     int="gelu"				,						_A		:     Optional[Any]=0.1				,						_A		:     List[Any]=0.1				,						_A		:     List[str]=512				,						_A		:     List[str]=16				,						_A		:     Dict=2				,						_A		:     Optional[Any]=0.02				,						_A		:     Any=3				,						_A		:     Optional[int]=4				,						_A		:     Any=None				,						)		->   Union[str, Any]:






       """simple docstring"""







       lowercase :       Any       =			parent
       lowercase :       Dict       =			batch_size
       lowercase :       Optional[Any]       =			seq_length
       lowercase :       Dict       =			is_training
       lowercase :       List[str]       =			use_input_mask
       lowercase :       str       =			use_token_type_ids
       lowercase :       Union[str, Any]       =			use_labels
       lowercase :       str       =			vocab_size
       lowercase :       Optional[int]       =			hidden_size
       lowercase :       List[str]       =			num_hidden_layers
       lowercase :       Union[str, Any]       =			num_attention_heads
       lowercase :       Tuple       =			intermediate_size
       lowercase :       Optional[Any]       =			hidden_act
       lowercase :       List[str]       =			hidden_dropout_prob
       lowercase :       Any       =			attention_probs_dropout_prob
       lowercase :       str       =			max_position_embeddings
       lowercase :       str       =			type_vocab_size
       lowercase :       Tuple       =			type_sequence_label_size
       lowercase :       List[str]       =			initializer_range
       lowercase :       Tuple       =			num_labels
       lowercase :       str       =			num_choices
       lowercase :       Optional[int]       =			scope






  def  __a      (  self		:     List[str]				)		->   Union[str, Any]:






       """simple docstring"""







       lowercase :       int       =			ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size				)

       lowercase :       str       =			None
       if self.use_input_mask:
            lowercase :       str       =			random_attention_mask([self.batch_size, self.seq_length]				)

       lowercase :       Union[str, Any]       =			None
       if self.use_token_type_ids:
            lowercase :       Any       =			ids_tensor([self.batch_size, self.seq_length]				,						self.type_vocab_size				)

       lowercase :       Optional[int]       =			None
       lowercase :       List[Any]       =			None
       lowercase :       Any       =			None
       if self.use_labels:
            lowercase :       Dict       =			ids_tensor([self.batch_size]				,						self.type_sequence_label_size				)
            lowercase :       Tuple       =			ids_tensor([self.batch_size, self.seq_length]				,						self.num_labels				)
            lowercase :       int       =			ids_tensor([self.batch_size]				,						self.num_choices				)

       lowercase :       Dict       =			self.get_config()

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






  def  __a      (  self		:     List[Any]				)		->   str:






       """simple docstring"""







       return LlamaConfig(
           vocab_size=self.vocab_size				,						hidden_size=self.hidden_size				,						num_hidden_layers=self.num_hidden_layers				,						num_attention_heads=self.num_attention_heads				,						intermediate_size=self.intermediate_size				,						hidden_act=self.hidden_act				,						hidden_dropout_prob=self.hidden_dropout_prob				,						attention_probs_dropout_prob=self.attention_probs_dropout_prob				,						max_position_embeddings=self.max_position_embeddings				,						type_vocab_size=self.type_vocab_size				,						is_decoder=_A				,						initializer_range=self.initializer_range				,						)






  def  __a      (  self		:     List[Any]				,						_A		:     Dict				,						_A		:     List[Any]				,						_A		:     Dict				,						_A		:     Optional[int]				,						_A		:     Optional[Any]				,						_A		:     Union[str, Any]				,						_A		:     Optional[Any]				)		->   Optional[int]:






       """simple docstring"""







       lowercase :       str       =			LlamaModel(config=_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Union[str, Any]       =			model(_A				,						attention_mask=_A				)
       lowercase :       Union[str, Any]       =			model(_A				)
       self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.hidden_size)				)






  def  __a      (  self		:     int				,						_A		:     Optional[int]				,						_A		:     Dict				,						_A		:     Optional[int]				,						_A		:     int				,						_A		:     str				,						_A		:     str				,						_A		:     List[Any]				,						_A		:     int				,						_A		:     List[Any]				,						)		->   List[str]:






       """simple docstring"""







       lowercase :       List[str]       =			True
       lowercase :       str       =			LlamaModel(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Any       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						)
       lowercase :       int       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						)
       lowercase :       Dict       =			model(_A				,						attention_mask=_A				)
       self.parent.assertEqual(result.last_hidden_state.shape				,						(self.batch_size, self.seq_length, self.hidden_size)				)






  def  __a      (  self		:     List[str]				,						_A		:     List[Any]				,						_A		:     Tuple				,						_A		:     Optional[Any]				,						_A		:     Optional[int]				,						_A		:     Dict				,						_A		:     Optional[Any]				,						_A		:     List[str]				,						_A		:     Optional[int]				,						_A		:     int				,						)		->   Tuple:






       """simple docstring"""







       lowercase :       Dict       =			LlamaForCausalLM(config=_A				)
       model.to(_A				)
       model.eval()
       lowercase :       str       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)				)






  def  __a      (  self		:     Union[str, Any]				,						_A		:     List[str]				,						_A		:     List[Any]				,						_A		:     Dict				,						_A		:     Optional[int]				,						_A		:     Union[str, Any]				,						_A		:     Optional[Any]				,						_A		:     Any				,						_A		:     Any				,						_A		:     Dict				,						)		->   Dict:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			True
       lowercase :       List[Any]       =			True
       lowercase :       Union[str, Any]       =			LlamaForCausalLM(config=_A				)
       model.to(_A				)
       model.eval()

       # first forward pass
       lowercase :       List[str]       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						use_cache=_A				,						)
       lowercase :       Dict       =			outputs.past_key_values

       # create hypothetical multiple next token and extent to next_input_ids
       lowercase :       Union[str, Any]       =			ids_tensor((self.batch_size, 3)				,						config.vocab_size				)
       lowercase :       List[str]       =			ids_tensor((self.batch_size, 3)				,						vocab_size=2				)

       # append to next input_ids and
       lowercase :       Union[str, Any]       =			torch.cat([input_ids, next_tokens]				,						dim=-1				)
       lowercase :       Any       =			torch.cat([input_mask, next_mask]				,						dim=-1				)

       lowercase :       Union[str, Any]       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						output_hidden_states=_A				,						)['''hidden_states'''][0]
       lowercase :       List[Any]       =			model(
           _A				,						attention_mask=_A				,						encoder_hidden_states=_A				,						encoder_attention_mask=_A				,						past_key_values=_A				,						output_hidden_states=_A				,						)['''hidden_states'''][0]

       # select random slice
       lowercase :       Optional[Any]       =			ids_tensor((1,)				,						output_from_past.shape[-1]				).item()
       lowercase :       Optional[Any]       =			output_from_no_past[:, -3:, random_slice_idx].detach()
       lowercase :       Dict       =			output_from_past[:, :, random_slice_idx].detach()

       self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]				)

       # test that outputs are equal for slice
       self.parent.assertTrue(torch.allclose(_A				,						_A				,						atol=1E-3				)				)






  def  __a      (  self		:     Optional[Any]				)		->   List[str]:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			self.prepare_config_and_inputs()
       (
           lowercase	
       ) :       Any       =			config_and_inputs
       lowercase :       Union[str, Any]       =			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict


@require_torch
class 						_A (  _lowerCamelCase      ,       _lowerCamelCase      ,       _lowerCamelCase      ,       unittest.TestCase				):
  _UpperCamelCase       :      str			=      (LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
  _UpperCamelCase       :      Any			=      (LlamaForCausalLM,) if is_torch_available() else ()
  _UpperCamelCase       :      Dict			=      (
      {
          '''feature-extraction''': LlamaModel,
          '''text-classification''': LlamaForSequenceClassification,
          '''text-generation''': LlamaForCausalLM,
          '''zero-shot''': LlamaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  _UpperCamelCase       :      str			=      False
  _UpperCamelCase       :      Optional[int]			=      False






  def  __a      (  self		:     List[str]				)		->   int:






       """simple docstring"""







       lowercase :       Optional[Any]       =			LlamaModelTester(self				)
       lowercase :       Optional[int]       =			ConfigTester(self				,						config_class=_A				,						hidden_size=37				)






  def  __a      (  self		:     List[str]				)		->   Dict:






       """simple docstring"""







       self.config_tester.run_common_tests()






  def  __a      (  self		:     Tuple				)		->   List[str]:






       """simple docstring"""







       lowercase :       Dict       =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_A				)






  def  __a      (  self		:     Dict				)		->   str:






       """simple docstring"""







       lowercase :       str       =			self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
            lowercase :       List[Any]       =			type
            self.model_tester.create_and_check_model(*_A				)






  def  __a      (  self		:     str				)		->   Union[str, Any]:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       Union[str, Any]       =			3
       lowercase :       Union[str, Any]       =			input_dict['''input_ids''']
       lowercase :       Any       =			input_ids.ne(1				).to(_A				)
       lowercase :       List[str]       =			ids_tensor([self.model_tester.batch_size]				,						self.model_tester.type_sequence_label_size				)
       lowercase :       Optional[int]       =			LlamaForSequenceClassification(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Dict       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.assertEqual(result.logits.shape				,						(self.model_tester.batch_size, self.model_tester.num_labels)				)






  def  __a      (  self		:     Any				)		->   int:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       Tuple       =			3
       lowercase :       List[str]       =			'''single_label_classification'''
       lowercase :       Union[str, Any]       =			input_dict['''input_ids''']
       lowercase :       int       =			input_ids.ne(1				).to(_A				)
       lowercase :       Union[str, Any]       =			ids_tensor([self.model_tester.batch_size]				,						self.model_tester.type_sequence_label_size				)
       lowercase :       Union[str, Any]       =			LlamaForSequenceClassification(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Tuple       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.assertEqual(result.logits.shape				,						(self.model_tester.batch_size, self.model_tester.num_labels)				)






  def  __a      (  self		:     Optional[Any]				)		->   int:






       """simple docstring"""







       lowercase :       List[Any]       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       Optional[Any]       =			3
       lowercase :       Dict       =			'''multi_label_classification'''
       lowercase :       int       =			input_dict['''input_ids''']
       lowercase :       str       =			input_ids.ne(1				).to(_A				)
       lowercase :       Any       =			ids_tensor(
           [self.model_tester.batch_size, config.num_labels]				,						self.model_tester.type_sequence_label_size				).to(torch.float				)
       lowercase :       Optional[Any]       =			LlamaForSequenceClassification(_A				)
       model.to(_A				)
       model.eval()
       lowercase :       Tuple       =			model(_A				,						attention_mask=_A				,						labels=_A				)
       self.assertEqual(result.logits.shape				,						(self.model_tester.batch_size, self.model_tester.num_labels)				)






  @unittest.skip('''LLaMA buffers include complex numbers, which breaks this test'''				)
  def  __a      (  self		:     Union[str, Any]				)		->   str:






       """simple docstring"""







       pass






  @parameterized.expand([('''linear''',), ('''dynamic''',)]				)
  def  __a      (  self		:     int				,						_A		:     List[str]				)		->   Dict:






       """simple docstring"""







       lowercase :       Tuple       =			self.model_tester.prepare_config_and_inputs_for_common()
       lowercase :       int       =			ids_tensor([1, 10]				,						config.vocab_size				)
       lowercase :       Dict       =			ids_tensor([1, int(config.max_position_embeddings * 1.5				)]				,						config.vocab_size				)

       set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
       lowercase :       Dict       =			LlamaModel(_A				)
       original_model.to(_A				)
       original_model.eval()
       lowercase :       Tuple       =			original_model(_A				).last_hidden_state
       lowercase :       Optional[Any]       =			original_model(_A				).last_hidden_state

       set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
       lowercase :       List[Any]       =			{'''type''': scaling_type, '''factor''': 10.0}
       lowercase :       Union[str, Any]       =			LlamaModel(_A				)
       scaled_model.to(_A				)
       scaled_model.eval()
       lowercase :       Union[str, Any]       =			scaled_model(_A				).last_hidden_state
       lowercase :       int       =			scaled_model(_A				).last_hidden_state

       # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
       # maximum sequence length, so the outputs for the short input should match.
       if scaling_type == "dynamic":
            self.assertTrue(torch.allclose(_A				,						_A				,						atol=1E-5				)				)
       else:
            self.assertFalse(torch.allclose(_A				,						_A				,						atol=1E-5				)				)

       # The output should be different for long inputs
       self.assertFalse(torch.allclose(_A				,						_A				,						atol=1E-5				)				)


@require_torch
class 						_A (  unittest.TestCase				):






  @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''				)
  @slow
  def  __a      (  self		:     int				)		->   List[str]:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       List[str]       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-7b-hf'''				,						device_map='''auto'''				)
       lowercase :       Union[str, Any]       =			model(torch.tensor([input_ids]				)				)
       # Expected mean on dim = -1
       lowercase :       Optional[int]       =			torch.tensor([[-6.6_550, -4.1_227, -4.9_859, -3.2_406, 0.8_262, -3.0_033, 1.2_964, -3.3_699]]				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       lowercase :       List[str]       =			torch.tensor([-12.8_281, -7.4_453, -0.4_639, -8.0_625, -7.2_500, -8.0_000, -6.4_883, -7.7_695, -7.8_438, -7.0_312, -6.2_188, -7.1_328, -1.8_496, 1.9_961, -8.6_250, -6.7_227, -12.8_281, -6.9_492, -7.0_742, -7.7_852, -7.5_820, -7.9_062, -6.9_375, -7.9_805, -8.3_438, -8.1_562, -8.0_469, -7.6_250, -7.7_422, -7.3_398,]				)
       # fmt: on
       torch.testing.assert_close(out[0, 0, :30]				,						_A				,						atol=1E-5				,						rtol=1E-5				)






  @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''				)
  @slow
  def  __a      (  self		:     List[str]				)		->   Any:






       """simple docstring"""







       lowercase :       int       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       List[str]       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-hf'''				,						device_map='''auto'''				)
       lowercase :       str       =			model(torch.tensor(_A				)				)
       # Expected mean on dim = -1
       lowercase :       Any       =			torch.tensor([[-2.0_622, -1.2_794, -1.1_638, -0.9_788, -1.4_603, -1.0_238, -1.7_893, -1.4_411]]				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       lowercase :       Optional[int]       =			torch.tensor([-8.1_406, -8.0_547, 2.7_461, -1.2_344, -0.1_448, -1.8_262, -1.0_020, -1.8_154, -1.6_895, -1.8_516, -2.3_574, -0.9_277, 3.7_598, 6.5_742, -1.2_998, -0.1_177, -8.1_406, -2.9_688, -2.9_199, -3.1_699, -3.5_254, -2.3_555, -2.7_988, -3.4_141, -2.8_262, -4.5_195, -3.3_379, -3.3_164, -2.7_832, -3.0_273]				)
       # fmt: on
       torch.testing.assert_close(out[0, 0, :30]				,						_A				,						atol=1E-5				,						rtol=1E-5				)






  @unittest.skip('''Logits are not exactly the same, once we fix the instabalities somehow, will update!'''				)
  @slow
  def  __a      (  self		:     int				)		->   Dict:






       """simple docstring"""







       lowercase :       str       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       int       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''				,						device_map='''auto'''				)
       lowercase :       str       =			model(torch.tensor(_A				)				)
       # Expected mean on dim = -1
       lowercase :       Optional[Any]       =			torch.tensor([[-0.8_562, -1.8_520, -0.7_551, -0.4_162, -1.5_161, -1.2_038, -2.4_823, -2.3_254]]				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # slicing logits[0, 0, 0:30]
       # fmt: off
       lowercase :       Any       =			torch.tensor([-2.2_227, 4.8_828, 0.9_023, -0.4_578, -0.7_871, -0.1_033, -0.6_221, -0.5_786, -0.7_803, -1.0_674, -1.2_920, -0.1_570, 0.8_008, 2.0_723, -0.9_497, 0.2_771, -2.2_227, -0.7_612, -1.4_346, -1.2_061, -1.6_426, -0.3_000, -0.7_139, -1.1_934, -1.8_691, -1.6_973, -1.5_947, -1.2_705, -0.3_523, -0.5_513]				)
       # fmt: on
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)






  @unittest.skip(
      '''Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test'''				)
  @slow
  def  __a      (  self		:     int				)		->   List[Any]:






       """simple docstring"""







       lowercase :       int       =			[1, 306, 4_658, 278, 6_593, 310, 2_834, 338]
       lowercase :       List[Any]       =			LlamaForCausalLM.from_pretrained('''meta-llama/Llama-2-70b-hf'''				,						device_map='''auto'''				)
       lowercase :       Any       =			model(torch.tensor(_A				)				)

       lowercase :       Union[str, Any]       =			torch.tensor(
           [[-4.2_327, -3.3_360, -4.6_665, -4.7_631, -1.8_180, -3.4_170, -1.4_211, -3.1_810]]				,						dtype=torch.floataa				)
       torch.testing.assert_close(out.mean(-1				)				,						_A				,						atol=1E-2				,						rtol=1E-2				)
       # fmt: off
       lowercase :       List[str]       =			torch.tensor([-9.4_922, -3.9_551, 1.7_998, -5.6_758, -5.1_055, -5.8_984, -4.8_320, -6.8_086, -6.5_391, -5.6_172, -5.5_820, -5.5_352, 1.7_881, 3.6_289, -6.5_117, -3.4_785, -9.5_000, -6.0_352, -6.8_125, -6.0_195, -6.6_836, -5.4_727, -6.2_812, -6.0_391, -7.3_398, -7.4_297, -7.4_844, -6.5_820, -5.8_789, -5.5_312]				)
       # fmt: on
       torch.testing.assert_close(out[0, 0, :30]				,						_A				,						atol=1E-5				,						rtol=1E-5				)






  @unittest.skip('''Model is curently gated'''				)
  @slow
  def  __a      (  self		:     int				)		->   Optional[int]:






       """simple docstring"""







       lowercase :       Any       =			'''Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the "princi'''
       lowercase :       str       =			'''Simply put, the theory of relativity states that '''
       lowercase :       Optional[Any]       =			LlamaTokenizer.from_pretrained('''meta-llama/Llama-2-13b-chat-hf'''				)
       lowercase :       Union[str, Any]       =			tokenizer.encode(_A				,						return_tensors='''pt'''				)
       lowercase :       Tuple       =			LlamaForCausalLM.from_pretrained(
           '''meta-llama/Llama-2-13b-chat-hf'''				,						device_map='''sequential'''				,						use_safetensors=_A				)

       # greedy generation outputs
       lowercase :       List[str]       =			model.generate(_A				,						max_new_tokens=64				,						top_p=_A				,						temperature=1				,						do_sample=_A				)
       lowercase :       List[Any]       =			tokenizer.decode(generated_ids[0]				,						skip_special_tokens=_A				)
       self.assertEqual(_A				,						_A				)





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_						=			None
lowerCAmelCase_						=			{
    '7B': 1_10_08,
    '13B': 1_38_24,
    '30B': 1_79_20,
    '65B': 2_20_16,
    '70B': 2_86_72,
}
lowerCAmelCase_						=			{
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}






def  snake_case(    __magic_name__			,		__magic_name__=1			,		__magic_name__=2_56   )			->							List[Any]:





     '''simple docstring'''


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






def  snake_case(    __magic_name__   )			->							Union[str, Any]:





     '''simple docstring'''


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






def  snake_case(    __magic_name__			,		__magic_name__   )			->							Optional[int]:





     '''simple docstring'''


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






def  snake_case(    __magic_name__			,		__magic_name__			,		__magic_name__			,		__magic_name__=True   )			->							Tuple:





     '''simple docstring'''


     os.makedirs(__magic_name__			,		exist_ok=__magic_name__   )
     lowercase :       List[Any]       =			os.path.join(__magic_name__			,		'''tmp'''   )
     os.makedirs(__magic_name__			,		exist_ok=__magic_name__   )

     lowercase :       List[Any]       =			read_json(os.path.join(__magic_name__			,		'''params.json'''   )   )
     lowercase :       int       =			NUM_SHARDS[model_size]
     lowercase :       str       =			params['''n_layers''']
     lowercase :       Optional[int]       =			params['''n_heads''']
     lowercase :       str       =			n_heads // num_shards
     lowercase :       Dict       =			params['''dim''']
     lowercase :       int       =			dim // n_heads
     lowercase :       List[str]       =			1_0_0_0_0.0
     lowercase :       Optional[Any]       =			1.0 / (base ** (torch.arange(0			,		__magic_name__			,		2   ).float() / dims_per_head))

     if "n_kv_heads" in params:
          lowercase :       Optional[Any]       =			params['''n_kv_heads''']  # for GQA / MQA
          lowercase :       Union[str, Any]       =			n_heads_per_shard // num_key_value_heads
          lowercase :       Any       =			dim // num_key_value_heads
     else:  # compatibility with other checkpoints
          lowercase :       Optional[Any]       =			n_heads
          lowercase :       str       =			n_heads_per_shard
          lowercase :       Any       =			dim

     # permute for sliced rotary
     def permute(__magic_name__			,		__magic_name__=n_heads			,		__magic_name__=dim			,		__magic_name__=dim   ):
          return w.view(__magic_name__			,		dima // n_heads // 2			,		2			,		__magic_name__   ).transpose(1			,		2   ).reshape(__magic_name__			,		__magic_name__   )

     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.)
          lowercase :       Tuple       =			torch.load(os.path.join(__magic_name__			,		'''consolidated.00.pth'''   )			,		map_location='''cpu'''   )
     else:
          # Sharded
          lowercase :       str       =			[
              torch.load(os.path.join(__magic_name__			,		F"""consolidated.{i:02d}.pth"""   )			,		map_location='''cpu'''   )
              for i in range(__magic_name__   )
          ]
     lowercase :       Tuple       =			0
     lowercase :       str       =			{'''weight_map''': {}}
     for layer_i in range(__magic_name__   ):
          lowercase :       List[Any]       =			F"""pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"""
          if model_size == "7B":
               # Unsharded
               lowercase :       int       =			{
                   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.

               lowercase :       List[Any]       =			{
                   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(),
               }
               lowercase :       Tuple       =			permute(
                   torch.cat(
                       [
                           loaded[i][F"""layers.{layer_i}.attention.wq.weight"""].view(__magic_name__			,		__magic_name__			,		__magic_name__   )
                           for i in range(__magic_name__   )
                       ]			,		dim=0			,		).reshape(__magic_name__			,		__magic_name__   )   )
               lowercase :       Any       =			permute(
                   torch.cat(
                       [
                           loaded[i][F"""layers.{layer_i}.attention.wk.weight"""].view(
                               __magic_name__			,		__magic_name__			,		__magic_name__   )
                           for i in range(__magic_name__   )
                       ]			,		dim=0			,		).reshape(__magic_name__			,		__magic_name__   )			,		__magic_name__			,		__magic_name__			,		__magic_name__			,		)
               lowercase :       Any       =			torch.cat(
                   [
                       loaded[i][F"""layers.{layer_i}.attention.wv.weight"""].view(
                           __magic_name__			,		__magic_name__			,		__magic_name__   )
                       for i in range(__magic_name__   )
                   ]			,		dim=0			,		).reshape(__magic_name__			,		__magic_name__   )

               lowercase :       int       =			torch.cat(
                   [loaded[i][F"""layers.{layer_i}.attention.wo.weight"""] for i in range(__magic_name__   )]			,		dim=1   )
               lowercase :       str       =			torch.cat(
                   [loaded[i][F"""layers.{layer_i}.feed_forward.w1.weight"""] for i in range(__magic_name__   )]			,		dim=0   )
               lowercase :       Any       =			torch.cat(
                   [loaded[i][F"""layers.{layer_i}.feed_forward.w2.weight"""] for i in range(__magic_name__   )]			,		dim=1   )
               lowercase :       List[Any]       =			torch.cat(
                   [loaded[i][F"""layers.{layer_i}.feed_forward.w3.weight"""] for i in range(__magic_name__   )]			,		dim=0   )

          lowercase :       List[Any]       =			inv_freq
          for k, v in state_dict.items():
               lowercase :       Tuple       =			filename
               param_count += v.numel()
          torch.save(__magic_name__			,		os.path.join(__magic_name__			,		__magic_name__   )   )

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

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

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

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

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






def  snake_case(    __magic_name__			,		__magic_name__   )			->							Optional[int]:





     '''simple docstring'''


     lowercase :       str       =			LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
     print(F"""Saving a {tokenizer_class.__name__} to {tokenizer_path}."""   )
     lowercase :       Tuple       =			tokenizer_class(__magic_name__   )
     tokenizer.save_pretrained(__magic_name__   )






def  snake_case(    )			->							Optional[Any]:





     '''simple docstring'''


     lowercase :       List[str]       =			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=__magic_name__			,		help='''Whether or not to save using `safetensors`.'''   )
     lowercase :       List[str]       =			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			,		)
     lowercase :       Any       =			os.path.join(args.input_dir			,		'''tokenizer.model'''   )
     write_tokenizer(args.output_dir			,		__magic_name__   )


if __name__ == "__main__":
 main()