


import itertools
import json
import os
import unittest

from transformers import AddedToken, RobertaTokenizer, RobertaTokenizerFast
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin
@require_tokenizers
class 						lowerCamelCase_							(					SCREAMING_SNAKE_CASE_  ,  unittest.TestCase					):



					'''simple docstring'''







					__UpperCAmelCase							  =					RobertaTokenizer
					__UpperCAmelCase							  =					RobertaTokenizerFast
					__UpperCAmelCase							  =					True
					__UpperCAmelCase							  =					{"cls_token": "<s>"}






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								super().setUp()

								# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
								__lowercase			       = [
								    """l""",
								    """o""",
								    """w""",
								    """e""",
								    """r""",
								    """s""",
								    """t""",
								    """i""",
								    """d""",
								    """n""",
								    """\u0120""",
								    """\u0120l""",
								    """\u0120n""",
								    """\u0120lo""",
								    """\u0120low""",
								    """er""",
								    """\u0120lowest""",
								    """\u0120newer""",
								    """\u0120wider""",
								    """<unk>""",
								]
								__lowercase			       = dict(zip(__snake_case			,							range(len(__snake_case    )    )    )    )
								__lowercase			       = ["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
								__lowercase			       = {"""unk_token""": """<unk>"""}

								__lowercase			       = os.path.join(self.tmpdirname			,							VOCAB_FILES_NAMES['''vocab_file''']    )
								__lowercase			       = os.path.join(self.tmpdirname			,							VOCAB_FILES_NAMES['''merges_file''']    )
								with open(self.vocab_file			,							'''w'''			,							encoding='''utf-8'''    ) as fp:
											fp.write(json.dumps(__snake_case    ) + '''\n'''    )
								with open(self.merges_file			,							'''w'''			,							encoding='''utf-8'''    ) as fp:
											fp.write('''\n'''.join(__snake_case    )    )






					def A   (					self			,							**snake_case_    ) ->      str:







								'''simple docstring'''

								kwargs.update(self.special_tokens_map    )
								return self.tokenizer_class.from_pretrained(self.tmpdirname			,							**__snake_case    )






					def A   (					self			,							**snake_case_    ) ->      Optional[Any]:







								'''simple docstring'''

								kwargs.update(self.special_tokens_map    )
								return RobertaTokenizerFast.from_pretrained(self.tmpdirname			,							**__snake_case    )






					def A   (					self			,							snake_case_    ) ->      List[Any]:







								'''simple docstring'''

								__lowercase			       = """lower newer"""
								__lowercase			       = """lower newer"""
								return input_text, output_text






					def A   (					self    ) ->      Optional[Any]:







								'''simple docstring'''

								__lowercase			       = self.tokenizer_class(self.vocab_file			,							self.merges_file			,							**self.special_tokens_map    )
								__lowercase			       = """lower newer"""
								__lowercase			       = ["""l""", """o""", """w""", """er""", """\u0120""", """n""", """e""", """w""", """er"""]
								__lowercase			       = tokenizer.tokenize(__snake_case    )  # , add_prefix_space=True)
								self.assertListEqual(__snake_case			,							__snake_case    )

								__lowercase			       = tokens + [tokenizer.unk_token]
								__lowercase			       = [0, 1, 2, 1_5, 1_0, 9, 3, 2, 1_5, 1_9]
								self.assertListEqual(tokenizer.convert_tokens_to_ids(__snake_case    )			,							__snake_case    )






					def A   (					self    ) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = self.get_tokenizer()

								self.assertListEqual(tokenizer.encode('''Hello world!'''			,							add_special_tokens=__snake_case    )			,							[0, 3_1_4_1_4, 2_3_2, 3_2_8, 2]    )
								self.assertListEqual(
								    tokenizer.encode('''Hello world! cécé herlolip 418'''			,							add_special_tokens=__snake_case    )			,							[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]			,							)






					@slow
					def A   (					self    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = self.tokenizer_class.from_pretrained('''roberta-base'''    )

								__lowercase			       = tokenizer.encode('''sequence builders'''			,							add_special_tokens=__snake_case    )
								__lowercase			       = tokenizer.encode('''multi-sequence build'''			,							add_special_tokens=__snake_case    )

								__lowercase			       = tokenizer.encode(
								    '''sequence builders'''			,							add_special_tokens=__snake_case			,							add_prefix_space=__snake_case    )
								__lowercase			       = tokenizer.encode(
								    '''sequence builders'''			,							'''multi-sequence build'''			,							add_special_tokens=__snake_case			,							add_prefix_space=__snake_case    )

								__lowercase			       = tokenizer.build_inputs_with_special_tokens(__snake_case    )
								__lowercase			       = tokenizer.build_inputs_with_special_tokens(__snake_case			,							__snake_case    )

								assert encoded_sentence == encoded_text_from_decode
								assert encoded_pair == encoded_pair_from_decode






					def A   (					self    ) ->      str:







								'''simple docstring'''

								__lowercase			       = self.get_tokenizer()

								__lowercase			       = """Encode this sequence."""
								__lowercase			       = tokenizer.byte_encoder[""" """.encode('''utf-8'''    )[0]]

								# Testing encoder arguments
								__lowercase			       = tokenizer.encode(__snake_case			,							add_special_tokens=__snake_case			,							add_prefix_space=__snake_case    )
								__lowercase			       = tokenizer.convert_ids_to_tokens(encoded[0]    )[0]
								self.assertNotEqual(__snake_case			,							__snake_case    )

								__lowercase			       = tokenizer.encode(__snake_case			,							add_special_tokens=__snake_case			,							add_prefix_space=__snake_case    )
								__lowercase			       = tokenizer.convert_ids_to_tokens(encoded[0]    )[0]
								self.assertEqual(__snake_case			,							__snake_case    )

								tokenizer.add_special_tokens({'''bos_token''': '''<s>'''}    )
								__lowercase			       = tokenizer.encode(__snake_case			,							add_special_tokens=__snake_case    )
								__lowercase			       = tokenizer.convert_ids_to_tokens(encoded[1]    )[0]
								self.assertNotEqual(__snake_case			,							__snake_case    )

								# Testing spaces after special tokens
								__lowercase			       = """<mask>"""
								tokenizer.add_special_tokens(
								    {'''mask_token''': AddedToken(__snake_case			,							lstrip=__snake_case			,							rstrip=__snake_case    )}    )  # mask token has a left space
								__lowercase			       = tokenizer.convert_tokens_to_ids(__snake_case    )

								__lowercase			       = """Encode <mask> sequence"""
								__lowercase			       = """Encode <mask>sequence"""

								__lowercase			       = tokenizer.encode(__snake_case    )
								__lowercase			       = encoded.index(__snake_case    )
								__lowercase			       = tokenizer.convert_ids_to_tokens(encoded[mask_loc + 1]    )[0]
								self.assertEqual(__snake_case			,							__snake_case    )

								__lowercase			       = tokenizer.encode(__snake_case    )
								__lowercase			       = encoded.index(__snake_case    )
								__lowercase			       = tokenizer.convert_ids_to_tokens(encoded[mask_loc + 1]    )[0]
								self.assertNotEqual(__snake_case			,							__snake_case    )






					def A   (					self    ) ->      int:







								'''simple docstring'''

								pass






					def A   (					self    ) ->      Tuple:







								'''simple docstring'''

								for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
											with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'    ):
														__lowercase			       = self.rust_tokenizer_class.from_pretrained(__snake_case			,							**__snake_case    )
														__lowercase			       = self.tokenizer_class.from_pretrained(__snake_case			,							**__snake_case    )
														__lowercase			       = """A, <mask> AllenNLP sentence."""
														__lowercase			       = tokenizer_r.encode_plus(__snake_case			,							add_special_tokens=__snake_case			,							return_token_type_ids=__snake_case    )
														__lowercase			       = tokenizer_p.encode_plus(__snake_case			,							add_special_tokens=__snake_case			,							return_token_type_ids=__snake_case    )

														# token_type_ids should put 0 everywhere
														self.assertEqual(sum(tokens_r['''token_type_ids''']    )			,							sum(tokens_p['''token_type_ids''']    )    )

														# attention_mask should put 1 everywhere, so sum over length should be 1
														self.assertEqual(
														    sum(tokens_r['''attention_mask''']    ) / len(tokens_r['''attention_mask''']    )			,							sum(tokens_p['''attention_mask''']    ) / len(tokens_p['''attention_mask''']    )			,							)

														__lowercase			       = tokenizer_r.convert_ids_to_tokens(tokens_r['''input_ids''']    )
														__lowercase			       = tokenizer_p.convert_ids_to_tokens(tokens_p['''input_ids''']    )

														# Rust correctly handles the space before the mask while python doesnt
														self.assertSequenceEqual(tokens_p['''input_ids''']			,							[0, 2_5_0, 6, 5_0_2_6_4, 3_8_2_3, 4_8_7, 2_1_9_9_2, 3_6_4_5, 4, 2]    )
														self.assertSequenceEqual(tokens_r['''input_ids''']			,							[0, 2_5_0, 6, 5_0_2_6_4, 3_8_2_3, 4_8_7, 2_1_9_9_2, 3_6_4_5, 4, 2]    )

														self.assertSequenceEqual(
														    __snake_case			,							['''<s>''', '''A''', ''',''', '''<mask>''', '''ĠAllen''', '''N''', '''LP''', '''Ġsentence''', '''.''', '''</s>''']    )
														self.assertSequenceEqual(
														    __snake_case			,							['''<s>''', '''A''', ''',''', '''<mask>''', '''ĠAllen''', '''N''', '''LP''', '''Ġsentence''', '''.''', '''</s>''']    )






					def A   (					self    ) ->      str:







								'''simple docstring'''

								for trim_offsets, add_prefix_space in itertools.product([True, False]			,							repeat=2    ):
											__lowercase			       = self.rust_tokenizer_class.from_pretrained(
											    self.tmpdirname			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )

											__lowercase			       = json.loads(tokenizer_r.backend_tokenizer.pre_tokenizer.__getstate__()    )
											__lowercase			       = json.loads(tokenizer_r.backend_tokenizer.post_processor.__getstate__()    )

											self.assertEqual(pre_tokenizer_state['''add_prefix_space''']			,							__snake_case    )

											self.assertEqual(post_processor_state['''add_prefix_space''']			,							__snake_case    )
											self.assertEqual(post_processor_state['''trim_offsets''']			,							__snake_case    )






					def A   (					self    ) ->      Tuple:







								'''simple docstring'''

								for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
											with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'    ):
														__lowercase			       = """hello"""  # `hello` is a token in the vocabulary of `pretrained_name`
														__lowercase			       = F'{text_of_1_token} {text_of_1_token}'

														__lowercase			       = self.rust_tokenizer_class.from_pretrained(
														    __snake_case			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )
														__lowercase			       = tokenizer_r(__snake_case			,							return_offsets_mapping=__snake_case			,							add_special_tokens=__snake_case    )
														self.assertEqual(encoding.offset_mapping[0]			,							(0, len(__snake_case    ))    )
														self.assertEqual(
														    encoding.offset_mapping[1]			,							(len(__snake_case    ) + 1, len(__snake_case    ) + 1 + len(__snake_case    ))			,							)

														__lowercase			       = self.rust_tokenizer_class.from_pretrained(
														    __snake_case			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )
														__lowercase			       = tokenizer_r(__snake_case			,							return_offsets_mapping=__snake_case			,							add_special_tokens=__snake_case    )
														self.assertEqual(encoding.offset_mapping[0]			,							(0, len(__snake_case    ))    )
														self.assertEqual(
														    encoding.offset_mapping[1]			,							(len(__snake_case    ) + 1, len(__snake_case    ) + 1 + len(__snake_case    ))			,							)

														__lowercase			       = self.rust_tokenizer_class.from_pretrained(
														    __snake_case			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )
														__lowercase			       = tokenizer_r(__snake_case			,							return_offsets_mapping=__snake_case			,							add_special_tokens=__snake_case    )
														self.assertEqual(encoding.offset_mapping[0]			,							(0, len(__snake_case    ))    )
														self.assertEqual(
														    encoding.offset_mapping[1]			,							(len(__snake_case    ), len(__snake_case    ) + 1 + len(__snake_case    ))			,							)

														__lowercase			       = self.rust_tokenizer_class.from_pretrained(
														    __snake_case			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )
														__lowercase			       = tokenizer_r(__snake_case			,							return_offsets_mapping=__snake_case			,							add_special_tokens=__snake_case    )
														self.assertEqual(encoding.offset_mapping[0]			,							(0, len(__snake_case    ))    )
														self.assertEqual(
														    encoding.offset_mapping[1]			,							(len(__snake_case    ), len(__snake_case    ) + 1 + len(__snake_case    ))			,							)

														__lowercase			       = F' {text}'

														# tokenizer_r = self.rust_tokenizer_class.from_pretrained(
														#     pretrained_name, use_fast=True, add_prefix_space=True, trim_offsets=True
														# )
														# encoding = tokenizer_r(text, return_offsets_mapping=True, add_special_tokens=False)
														# self.assertEqual(encoding.offset_mapping[0], (1, 1 + len(text_of_1_token)))
														# self.assertEqual(
														#     encoding.offset_mapping[1],
														#     (1 + len(text_of_1_token) + 1, 1 + len(text_of_1_token) + 1 + len(text_of_1_token)),
														# )

														__lowercase			       = self.rust_tokenizer_class.from_pretrained(
														    __snake_case			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )
														__lowercase			       = tokenizer_r(__snake_case			,							return_offsets_mapping=__snake_case			,							add_special_tokens=__snake_case    )
														self.assertEqual(encoding.offset_mapping[0]			,							(1, 1 + len(__snake_case    ))    )
														self.assertEqual(
														    encoding.offset_mapping[1]			,							(1 + len(__snake_case    ) + 1, 1 + len(__snake_case    ) + 1 + len(__snake_case    ))			,							)

														__lowercase			       = self.rust_tokenizer_class.from_pretrained(
														    __snake_case			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )
														__lowercase			       = tokenizer_r(__snake_case			,							return_offsets_mapping=__snake_case			,							add_special_tokens=__snake_case    )
														self.assertEqual(encoding.offset_mapping[0]			,							(0, 1 + len(__snake_case    ))    )
														self.assertEqual(
														    encoding.offset_mapping[1]			,							(1 + len(__snake_case    ), 1 + len(__snake_case    ) + 1 + len(__snake_case    ))			,							)

														__lowercase			       = self.rust_tokenizer_class.from_pretrained(
														    __snake_case			,							use_fast=__snake_case			,							add_prefix_space=__snake_case			,							trim_offsets=__snake_case    )
														__lowercase			       = tokenizer_r(__snake_case			,							return_offsets_mapping=__snake_case			,							add_special_tokens=__snake_case    )
														self.assertEqual(encoding.offset_mapping[0]			,							(0, 1 + len(__snake_case    ))    )
														self.assertEqual(
														    encoding.offset_mapping[1]			,							(1 + len(__snake_case    ), 1 + len(__snake_case    ) + 1 + len(__snake_case    ))			,							)










"""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'
							)
							__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   a    (     __UpperCAmelCase      :    List[str]      ,   __UpperCAmelCase      :    Optional[Any]=1      ,   __UpperCAmelCase      :    Optional[int]=2_5_6  )    ->							List[str]:
						return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3  )  ) + multiple_of - 1) // multiple_of)



def   a    (     __UpperCAmelCase      :    Union[str, Any]  )    ->							Optional[Any]:
						with open(__UpperCAmelCase      ,   """r"""  ) as f:
												return json.load(__UpperCAmelCase  )



def   a    (     __UpperCAmelCase      :    Tuple      ,   __UpperCAmelCase      :    Union[str, Any]  )    ->							Optional[int]:
						with open(__UpperCAmelCase      ,   """w"""  ) as f:
												json.dump(__UpperCAmelCase      ,   __UpperCAmelCase  )



def   a    (     __UpperCAmelCase      :    List[Any]      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    str      ,   __UpperCAmelCase      :    Tuple=True  )    ->							Tuple:
						os.makedirs(__UpperCAmelCase      ,   exist_ok=__UpperCAmelCase  )
						__magic_name__:	Any							=       os.path.join(__UpperCAmelCase      ,   """tmp"""  )
						os.makedirs(__UpperCAmelCase      ,   exist_ok=__UpperCAmelCase  )

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

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

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

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

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

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

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

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

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

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

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

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



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



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


if __name__ == "__main__":
							main()



