



'''simple docstring'''






import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
						from transformers import SpeechaTextFeatureExtractor

UpperCamelCase__:   int								=			random.Random()


def  snake_case_     (					_lowerCAmelCase  :						List[Any]   ,      _lowerCAmelCase  :						int=1.0   ,      _lowerCAmelCase  :						Optional[Any]=None   ,      _lowerCAmelCase  :						int=None  )      ->   Any:
							if rng is None:
														UpperCAmelCase :       str 							=		global_rng

							UpperCAmelCase :       str 							=		[]
							for batch_idx in range(shape[0]  ):
														values.append([]  )
														for _ in range(shape[1]  ):
																					values[-1].append(rng.random() * scale  )

							return values

@require_torch
@require_torchaudio
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	def __init__(						self			:     int		,			__snake_case			:     str		,			__snake_case			:     List[Any]=7		,			__snake_case			:     Optional[int]=400		,			__snake_case			:     List[str]=2000		,			__snake_case			:     List[Any]=24		,			__snake_case			:     Optional[Any]=24		,			__snake_case			:     Dict=0.0		,			__snake_case			:     int=16000		,			__snake_case			:     Any=True		,			__snake_case			:     Dict=True		,			)			->		Tuple:
								UpperCAmelCase :       Tuple 							=		parent
								UpperCAmelCase :       int 							=		batch_size
								UpperCAmelCase :       int 							=		min_seq_length
								UpperCAmelCase :       Optional[int] 							=		max_seq_length
								UpperCAmelCase :       List[Any] 							=		(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
								UpperCAmelCase :       Optional[Any] 							=		feature_size
								UpperCAmelCase :       Union[str, Any] 							=		num_mel_bins
								UpperCAmelCase :       List[Any] 							=		padding_value
								UpperCAmelCase :       Optional[Any] 							=		sampling_rate
								UpperCAmelCase :       Union[str, Any] 							=		return_attention_mask
								UpperCAmelCase :       str 							=		do_normalize







	def      A	(						self			:     Dict		)			->		Dict:
								return {
								    "feature_size": self.feature_size,
								    "num_mel_bins": self.num_mel_bins,
								    "padding_value": self.padding_value,
								    "sampling_rate": self.sampling_rate,
								    "return_attention_mask": self.return_attention_mask,
								    "do_normalize": self.do_normalize,
								}







	def      A	(						self			:     Dict		,			__snake_case			:     List[Any]=False		,			__snake_case			:     int=False		)			->		List[str]:
								def _flatten(__snake_case			:     List[Any]		):
															return list(itertools.chain(*__snake_case		)		)

								if equal_length:
															UpperCAmelCase :       Union[str, Any] 							=		[floats_list((self.max_seq_length, self.feature_size)		) for _ in range(self.batch_size		)]
								else:
															# make sure that inputs increase in size
															UpperCAmelCase :       Optional[Any] 							=		[
															    floats_list((x, self.feature_size)		)
															    for x in range(self.min_seq_length		,			self.max_seq_length		,			self.seq_length_diff		)
															]
								if numpify:
															UpperCAmelCase :       List[Any] 							=		[np.asarray(__snake_case		) for x in speech_inputs]
								return speech_inputs

@require_torch
@require_torchaudio
class 				SCREAMING_SNAKE_CASE(							A__							,		unittest.TestCase			):







	"""simple docstring"""






	lowerCamelCase__			 =       SpeechaTextFeatureExtractor if is_speech_available() else None







	def      A	(						self			:     Union[str, Any]		)			->		Optional[Any]:
								UpperCAmelCase :       Optional[Any] 							=		SpeechaTextFeatureExtractionTester(self		)







	def      A	(						self			:     str		,			__snake_case			:     List[Any]		)			->		Dict:
								self.assertTrue(np.all(np.mean(__snake_case		,			axis=0		) < 1E-3		)		)
								self.assertTrue(np.all(np.abs(np.var(__snake_case		,			axis=0		) - 1		) < 1E-3		)		)







	def      A	(						self			:     Union[str, Any]		)			->		int:
								# Tests that all call wrap to encode_plus and batch_encode_plus
								UpperCAmelCase :       int 							=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
								# create three inputs of length 800, 1000, and 1200
								UpperCAmelCase :       Dict 							=		[floats_list((1, x)		)[0] for x in range(800		,			1400		,			200		)]
								UpperCAmelCase :       Union[str, Any] 							=		[np.asarray(__snake_case		) for speech_input in speech_inputs]

								# Test feature size
								UpperCAmelCase :       Optional[Any] 							=		feature_extractor(__snake_case		,			padding=__snake_case		,			return_tensors='''np'''		).input_features
								self.assertTrue(input_features.ndim == 3		)
								self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size		)

								# Test not batched input
								UpperCAmelCase :       Optional[int] 							=		feature_extractor(speech_inputs[0]		,			return_tensors='''np'''		).input_features
								UpperCAmelCase :       List[str] 							=		feature_extractor(np_speech_inputs[0]		,			return_tensors='''np'''		).input_features
								self.assertTrue(np.allclose(__snake_case		,			__snake_case		,			atol=1E-3		)		)

								# Test batched
								UpperCAmelCase :       Any 							=		feature_extractor(__snake_case		,			return_tensors='''np'''		).input_features
								UpperCAmelCase :       Union[str, Any] 							=		feature_extractor(__snake_case		,			return_tensors='''np'''		).input_features
								for enc_seq_a, enc_seq_a in zip(__snake_case		,			__snake_case		):
															self.assertTrue(np.allclose(__snake_case		,			__snake_case		,			atol=1E-3		)		)

								# Test 2-D numpy arrays are batched.
								UpperCAmelCase :       List[Any] 							=		[floats_list((1, x)		)[0] for x in (800, 800, 800)]
								UpperCAmelCase :       Optional[Any] 							=		np.asarray(__snake_case		)
								UpperCAmelCase :       Tuple 							=		feature_extractor(__snake_case		,			return_tensors='''np'''		).input_features
								UpperCAmelCase :       Tuple 							=		feature_extractor(__snake_case		,			return_tensors='''np'''		).input_features
								for enc_seq_a, enc_seq_a in zip(__snake_case		,			__snake_case		):
															self.assertTrue(np.allclose(__snake_case		,			__snake_case		,			atol=1E-3		)		)







	def      A	(						self			:     Optional[int]		)			->		Any:
								UpperCAmelCase :       Dict 							=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
								UpperCAmelCase :       Any 							=		[floats_list((1, x)		)[0] for x in range(800		,			1400		,			200		)]

								UpperCAmelCase :       List[str] 							=		['''longest''', '''max_length''', '''do_not_pad''']
								UpperCAmelCase :       Dict 							=		[None, 16, None]
								for max_length, padding in zip(__snake_case		,			__snake_case		):
															UpperCAmelCase :       Optional[Any] 							=		feature_extractor(
															    __snake_case		,			padding=__snake_case		,			max_length=__snake_case		,			return_attention_mask=__snake_case		)
															UpperCAmelCase :       Optional[Any] 							=		inputs.input_features
															UpperCAmelCase :       Dict 							=		inputs.attention_mask
															UpperCAmelCase :       Optional[int] 							=		[np.sum(__snake_case		) for x in attention_mask]

															self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]		)
															self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]		)
															self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]		)







	def      A	(						self			:     Optional[Any]		)			->		List[str]:
								UpperCAmelCase :       int 							=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
								UpperCAmelCase :       Tuple 							=		[floats_list((1, x)		)[0] for x in range(800		,			1400		,			200		)]

								UpperCAmelCase :       Optional[int] 							=		['''longest''', '''max_length''', '''do_not_pad''']
								UpperCAmelCase :       List[Any] 							=		[None, 16, None]
								for max_length, padding in zip(__snake_case		,			__snake_case		):
															UpperCAmelCase :       Tuple 							=		feature_extractor(
															    __snake_case		,			max_length=__snake_case		,			padding=__snake_case		,			return_tensors='''np'''		,			return_attention_mask=__snake_case		)
															UpperCAmelCase :       Optional[int] 							=		inputs.input_features
															UpperCAmelCase :       int 							=		inputs.attention_mask
															UpperCAmelCase :       List[str] 							=		[np.sum(__snake_case		) for x in attention_mask]

															self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]		)
															self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6		)
															self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]		)
															self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6		)
															self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]		)







	def      A	(						self			:     List[str]		)			->		Any:
								UpperCAmelCase :       Any 							=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
								UpperCAmelCase :       List[Any] 							=		[floats_list((1, x)		)[0] for x in range(800		,			1400		,			200		)]
								UpperCAmelCase :       Union[str, Any] 							=		feature_extractor(
								    __snake_case		,			padding='''max_length'''		,			max_length=4		,			truncation=__snake_case		,			return_tensors='''np'''		,			return_attention_mask=__snake_case		,			)
								UpperCAmelCase :       Union[str, Any] 							=		inputs.input_features
								UpperCAmelCase :       int 							=		inputs.attention_mask
								UpperCAmelCase :       str 							=		np.sum(attention_mask == 1		,			axis=1		)

								self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
								self._check_zero_mean_unit_variance(input_features[1]		)
								self._check_zero_mean_unit_variance(input_features[2]		)







	def      A	(						self			:     int		)			->		int:
								UpperCAmelCase :       int 							=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
								UpperCAmelCase :       Any 							=		[floats_list((1, x)		)[0] for x in range(800		,			1400		,			200		)]
								UpperCAmelCase :       List[Any] 							=		feature_extractor(
								    __snake_case		,			padding='''longest'''		,			max_length=4		,			truncation=__snake_case		,			return_tensors='''np'''		,			return_attention_mask=__snake_case		,			)
								UpperCAmelCase :       List[str] 							=		inputs.input_features
								UpperCAmelCase :       Optional[int] 							=		inputs.attention_mask
								UpperCAmelCase :       Dict 							=		np.sum(attention_mask == 1		,			axis=1		)

								self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
								self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]		)
								self._check_zero_mean_unit_variance(input_features[2]		)

								# make sure that if max_length < longest -> then pad to max_length
								self.assertEqual(input_features.shape		,			(3, 4, 24)		)

								UpperCAmelCase :       List[str] 							=		[floats_list((1, x)		)[0] for x in range(800		,			1400		,			200		)]
								UpperCAmelCase :       Any 							=		feature_extractor(
								    __snake_case		,			padding='''longest'''		,			max_length=16		,			truncation=__snake_case		,			return_tensors='''np'''		,			return_attention_mask=__snake_case		,			)
								UpperCAmelCase :       List[Any] 							=		inputs.input_features
								UpperCAmelCase :       str 							=		inputs.attention_mask
								UpperCAmelCase :       Any 							=		np.sum(attention_mask == 1		,			axis=1		)

								self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
								self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]		)
								self._check_zero_mean_unit_variance(input_features[2]		)

								# make sure that if max_length < longest -> then pad to max_length
								self.assertEqual(input_features.shape		,			(3, 6, 24)		)







	def      A	(						self			:     Any		)			->		Any:
								import torch

								UpperCAmelCase :       List[str] 							=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
								UpperCAmelCase :       str 							=		np.random.rand(100		,			32		).astype(np.floataa		)
								UpperCAmelCase :       int 							=		np_speech_inputs.tolist()

								for inputs in [py_speech_inputs, np_speech_inputs]:
															UpperCAmelCase :       Union[str, Any] 							=		feature_extractor.pad([{'''input_features''': inputs}]		,			return_tensors='''np'''		)
															self.assertTrue(np_processed.input_features.dtype == np.floataa		)
															UpperCAmelCase :       int 							=		feature_extractor.pad([{'''input_features''': inputs}]		,			return_tensors='''pt'''		)
															self.assertTrue(pt_processed.input_features.dtype == torch.floataa		)







	def      A	(						self			:     Optional[Any]		,			__snake_case			:     Optional[int]		)			->		Optional[Any]:
								from datasets import load_dataset

								UpperCAmelCase :       List[Any] 							=		load_dataset('''hf-internal-testing/librispeech_asr_dummy'''		,			'''clean'''		,			split='''validation'''		)
								# automatic decoding with librispeech
								UpperCAmelCase :       Union[str, Any] 							=		ds.sort('''id'''		).select(range(__snake_case		)		)[:num_samples]['''audio''']

								return [x["array"] for x in speech_samples]







	def      A	(						self			:     Dict		)			->		Optional[Any]:
								# fmt: off
								UpperCAmelCase :       Dict 							=		np.array([
								    -1.57_45, -1.77_13, -1.70_20, -1.60_69, -1.22_50, -1.11_05, -0.90_72, -0.82_41,
								    -1.23_10, -0.80_98, -0.33_20, -0.41_01, -0.79_85, -0.49_96, -0.82_13, -0.91_28,
								    -1.04_20, -1.12_86, -1.04_40, -0.79_99, -0.84_05, -1.22_75, -1.54_43, -1.46_25,
								]		)
								# fmt: on

								UpperCAmelCase :       List[str] 							=		self._load_datasamples(1		)
								UpperCAmelCase :       Any 							=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
								UpperCAmelCase :       Optional[int] 							=		feature_extractor(__snake_case		,			return_tensors='''pt'''		).input_features
								self.assertEquals(input_features.shape		,			(1, 584, 24)		)
								self.assertTrue(np.allclose(input_features[0, 0, :30]		,			__snake_case		,			atol=1E-4		)		)










'''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__:   str								=			None
UpperCamelCase__:   int								=			{
    "7B": 11008,
    "13B": 13824,
    "30B": 17920,
    "65B": 22016,
    "70B": 28672,
}
UpperCamelCase__:   List[Any]								=			{
    "7B": 1,
    "7Bf": 1,
    "13B": 2,
    "13Bf": 2,
    "30B": 4,
    "65B": 8,
    "70B": 8,
    "70Bf": 8,
}


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


def  snake_case_     (					_lowerCAmelCase  :						List[str]  )      ->   str:
							with open(_lowerCAmelCase   ,      '''r'''  ) as f:
														return json.load(_lowerCAmelCase  )


def  snake_case_     (					_lowerCAmelCase  :						Union[str, Any]   ,      _lowerCAmelCase  :						Any  )      ->   Optional[Any]:
							with open(_lowerCAmelCase   ,      '''w'''  ) as f:
														json.dump(_lowerCAmelCase   ,      _lowerCAmelCase  )


def  snake_case_     (					_lowerCAmelCase  :						Dict   ,      _lowerCAmelCase  :						Tuple   ,      _lowerCAmelCase  :						Any   ,      _lowerCAmelCase  :						Any=True  )      ->   List[Any]:
							os.makedirs(_lowerCAmelCase   ,      exist_ok=_lowerCAmelCase  )
							UpperCAmelCase :       int 							=		os.path.join(_lowerCAmelCase   ,      '''tmp'''  )
							os.makedirs(_lowerCAmelCase   ,      exist_ok=_lowerCAmelCase  )

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

							if "n_kv_heads" in params:
														UpperCAmelCase :       Tuple 							=		params['''n_kv_heads''']  # for GQA / MQA
														UpperCAmelCase :       Optional[int] 							=		n_heads_per_shard // num_key_value_heads
														UpperCAmelCase :       Optional[Any] 							=		dim // num_key_value_heads
							else:  # compatibility with other checkpoints
														UpperCAmelCase :       List[str] 							=		n_heads
														UpperCAmelCase :       Optional[int] 							=		n_heads_per_shard
														UpperCAmelCase :       List[str] 							=		dim

							# permute for sliced rotary
							def permute(_lowerCAmelCase  :						Tuple   ,      _lowerCAmelCase  :						Optional[Any]=n_heads   ,      _lowerCAmelCase  :						int=dim   ,      _lowerCAmelCase  :						Dict=dim  ):
														return w.view(_lowerCAmelCase   ,      dima // n_heads // 2   ,      2   ,      _lowerCAmelCase  ).transpose(1   ,      2  ).reshape(_lowerCAmelCase   ,      _lowerCAmelCase  )

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

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

																					UpperCAmelCase :       Optional[int] 							=		torch.cat(
																					    [loaded[i][f"""layers.{layer_i}.attention.wo.weight"""] for i in range(_lowerCAmelCase  )]   ,      dim=1  )
																					UpperCAmelCase :       Any 							=		torch.cat(
																					    [loaded[i][f"""layers.{layer_i}.feed_forward.w1.weight"""] for i in range(_lowerCAmelCase  )]   ,      dim=0  )
																					UpperCAmelCase :       str 							=		torch.cat(
																					    [loaded[i][f"""layers.{layer_i}.feed_forward.w2.weight"""] for i in range(_lowerCAmelCase  )]   ,      dim=1  )
																					UpperCAmelCase :       Tuple 							=		torch.cat(
																					    [loaded[i][f"""layers.{layer_i}.feed_forward.w3.weight"""] for i in range(_lowerCAmelCase  )]   ,      dim=0  )

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

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

							for k, v in state_dict.items():
														UpperCAmelCase :       Optional[int] 							=		filename
														param_count += v.numel()
							torch.save(_lowerCAmelCase   ,      os.path.join(_lowerCAmelCase   ,      _lowerCAmelCase  )  )

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

							# 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 :       Optional[int] 							=		LlamaForCausalLM.from_pretrained(_lowerCAmelCase   ,      torch_dtype=torch.floataa   ,      low_cpu_mem_usage=_lowerCAmelCase  )
							# Avoid saving this as part of the config.
							del model.config._name_or_path

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


def  snake_case_     (					_lowerCAmelCase  :						int   ,      _lowerCAmelCase  :						Union[str, Any]  )      ->   List[str]:
							# Initialize the tokenizer based on the `spm` model
							UpperCAmelCase :       Dict 							=		LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
							print(f"""Saving a {tokenizer_class.__name__} to {tokenizer_path}."""  )
							UpperCAmelCase :       List[Any] 							=		tokenizer_class(_lowerCAmelCase  )
							tokenizer.save_pretrained(_lowerCAmelCase  )


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


if __name__ == "__main__":
						main()





