



import os
import unittest

from transformers import BertTokenizerFast
from transformers.models.bert.tokenization_bert import (
    VOCAB_FILES_NAMES,
    BasicTokenizer,
    BertTokenizer,
    WordpieceTokenizer,
    _is_control,
    _is_punctuation,
    _is_whitespace,
)
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin, filter_non_english





@require_tokenizers
class 					__lowerCAmelCase (     lowerCamelCase__       ,						unittest.TestCase       ):
				__lowerCamelCase =				BertTokenizer
				__lowerCamelCase =				BertTokenizerFast
				__lowerCamelCase =				True
				__lowerCamelCase =				True
				__lowerCamelCase =				filter_non_english




				def 							snake_case   (							self		):



									"""simple docstring"""





									super().setUp()

									_lowerCAmelCase						=    [
									    """[UNK]""",
									    """[CLS]""",
									    """[SEP]""",
									    """[PAD]""",
									    """[MASK]""",
									    """want""",
									    """##want""",
									    """##ed""",
									    """wa""",
									    """un""",
									    """runn""",
									    """##ing""",
									    """,""",
									    """low""",
									    """lowest""",
									]
									_lowerCAmelCase						=    os.path.join(self.tmpdirname						,				VOCAB_FILES_NAMES["""vocab_file"""]		)
									with open(self.vocab_file						,				"""w"""						,				encoding="""utf-8"""		) as vocab_writer:
														vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]		)		)




				def 							snake_case   (							self						,				_snake_case		):



									"""simple docstring"""





									_lowerCAmelCase						=    """UNwant\u00E9d,running"""
									_lowerCAmelCase						=    """unwanted, running"""
									return input_text, output_text




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.tokenizer_class(self.vocab_file		)

									_lowerCAmelCase						=    tokenizer.tokenize("""UNwant\u00E9d,running"""		)
									self.assertListEqual(_snake_case						,				["""un""", """##want""", """##ed""", """,""", """runn""", """##ing"""]		)
									self.assertListEqual(tokenizer.convert_tokens_to_ids(_snake_case		)						,				[9, 6, 7, 12, 10, 11]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									if not self.test_rust_tokenizer:
														return

									_lowerCAmelCase						=    self.get_tokenizer()
									_lowerCAmelCase						=    self.get_rust_tokenizer()

									_lowerCAmelCase						=    """UNwant\u00E9d,running"""

									_lowerCAmelCase						=    tokenizer.tokenize(_snake_case		)
									_lowerCAmelCase						=    rust_tokenizer.tokenize(_snake_case		)
									self.assertListEqual(_snake_case						,				_snake_case		)

									_lowerCAmelCase						=    tokenizer.encode(_snake_case						,				add_special_tokens=_snake_case		)
									_lowerCAmelCase						=    rust_tokenizer.encode(_snake_case						,				add_special_tokens=_snake_case		)
									self.assertListEqual(_snake_case						,				_snake_case		)

									_lowerCAmelCase						=    self.get_rust_tokenizer()
									_lowerCAmelCase						=    tokenizer.encode(_snake_case		)
									_lowerCAmelCase						=    rust_tokenizer.encode(_snake_case		)
									self.assertListEqual(_snake_case						,				_snake_case		)

									# With lower casing
									_lowerCAmelCase						=    self.get_tokenizer(do_lower_case=_snake_case		)
									_lowerCAmelCase						=    self.get_rust_tokenizer(do_lower_case=_snake_case		)

									_lowerCAmelCase						=    """UNwant\u00E9d,running"""

									_lowerCAmelCase						=    tokenizer.tokenize(_snake_case		)
									_lowerCAmelCase						=    rust_tokenizer.tokenize(_snake_case		)
									self.assertListEqual(_snake_case						,				_snake_case		)

									_lowerCAmelCase						=    tokenizer.encode(_snake_case						,				add_special_tokens=_snake_case		)
									_lowerCAmelCase						=    rust_tokenizer.encode(_snake_case						,				add_special_tokens=_snake_case		)
									self.assertListEqual(_snake_case						,				_snake_case		)

									_lowerCAmelCase						=    self.get_rust_tokenizer()
									_lowerCAmelCase						=    tokenizer.encode(_snake_case		)
									_lowerCAmelCase						=    rust_tokenizer.encode(_snake_case		)
									self.assertListEqual(_snake_case						,				_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer()

									self.assertListEqual(tokenizer.tokenize("""ah\u535A\u63A8zz"""		)						,				["""ah""", """\u535A""", """\u63A8""", """zz"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHeLLo!how  \n Are yoU?  """		)						,				["""hello""", """!""", """how""", """are""", """you""", """?"""]		)
									self.assertListEqual(tokenizer.tokenize("""H\u00E9llo"""		)						,				["""hello"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case						,				strip_accents=_snake_case		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHäLLo!how  \n Are yoU?  """		)						,				["""hällo""", """!""", """how""", """are""", """you""", """?"""]		)
									self.assertListEqual(tokenizer.tokenize("""H\u00E9llo"""		)						,				["""h\u00E9llo"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case						,				strip_accents=_snake_case		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHäLLo!how  \n Are yoU?  """		)						,				["""hallo""", """!""", """how""", """are""", """you""", """?"""]		)
									self.assertListEqual(tokenizer.tokenize("""H\u00E9llo"""		)						,				["""hello"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHäLLo!how  \n Are yoU?  """		)						,				["""hallo""", """!""", """how""", """are""", """you""", """?"""]		)
									self.assertListEqual(tokenizer.tokenize("""H\u00E9llo"""		)						,				["""hello"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHeLLo!how  \n Are yoU?  """		)						,				["""HeLLo""", """!""", """how""", """Are""", """yoU""", """?"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case						,				strip_accents=_snake_case		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHäLLo!how  \n Are yoU?  """		)						,				["""HäLLo""", """!""", """how""", """Are""", """yoU""", """?"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case						,				strip_accents=_snake_case		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHäLLo!how  \n Are yoU?  """		)						,				["""HaLLo""", """!""", """how""", """Are""", """yoU""", """?"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer(do_lower_case=_snake_case						,				never_split=["""[UNK]"""]		)

									self.assertListEqual(
									    tokenizer.tokenize(""" \tHeLLo!how  \n Are yoU? [UNK]"""		)						,				["""HeLLo""", """!""", """how""", """Are""", """yoU""", """?""", """[UNK]"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    BasicTokenizer()
									_lowerCAmelCase						=    """a\n'll !!to?'d of, can't."""
									_lowerCAmelCase						=    ["""a""", """'""", """ll""", """!""", """!""", """to""", """?""", """'""", """d""", """of""", """,""", """can""", """'""", """t""", """."""]
									self.assertListEqual(tokenizer.tokenize(_snake_case		)						,				_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    ["""[UNK]""", """[CLS]""", """[SEP]""", """want""", """##want""", """##ed""", """wa""", """un""", """runn""", """##ing"""]

									_lowerCAmelCase						=    {}
									for i, token in enumerate(_snake_case		):
														_lowerCAmelCase						=    i
									_lowerCAmelCase						=    WordpieceTokenizer(vocab=_snake_case						,				unk_token="""[UNK]"""		)

									self.assertListEqual(tokenizer.tokenize(""""""		)						,				[]		)

									self.assertListEqual(tokenizer.tokenize("""unwanted running"""		)						,				["""un""", """##want""", """##ed""", """runn""", """##ing"""]		)

									self.assertListEqual(tokenizer.tokenize("""unwantedX running"""		)						,				["""[UNK]""", """runn""", """##ing"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									self.assertTrue(_is_whitespace(""" """		)		)
									self.assertTrue(_is_whitespace("""\t"""		)		)
									self.assertTrue(_is_whitespace("""\r"""		)		)
									self.assertTrue(_is_whitespace("""\n"""		)		)
									self.assertTrue(_is_whitespace("""\u00A0"""		)		)

									self.assertFalse(_is_whitespace("""A"""		)		)
									self.assertFalse(_is_whitespace("""-"""		)		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									self.assertTrue(_is_control("""\u0005"""		)		)

									self.assertFalse(_is_control("""A"""		)		)
									self.assertFalse(_is_control(""" """		)		)
									self.assertFalse(_is_control("""\t"""		)		)
									self.assertFalse(_is_control("""\r"""		)		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									self.assertTrue(_is_punctuation("""-"""		)		)
									self.assertTrue(_is_punctuation("""$"""		)		)
									self.assertTrue(_is_punctuation("""`"""		)		)
									self.assertTrue(_is_punctuation("""."""		)		)

									self.assertFalse(_is_punctuation("""A"""		)		)
									self.assertFalse(_is_punctuation(""" """		)		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.get_tokenizer()
									_lowerCAmelCase						=    self.get_rust_tokenizer()

									# Example taken from the issue https://github.com/huggingface/tokenizers/issues/340
									self.assertListEqual([tokenizer.tokenize(_snake_case		) for t in ["""Test""", """\xad""", """test"""]]						,				[["""[UNK]"""], [], ["""[UNK]"""]]		)

									self.assertListEqual(
									    [rust_tokenizer.tokenize(_snake_case		) for t in ["""Test""", """\xad""", """test"""]]						,				[["""[UNK]"""], [], ["""[UNK]"""]]		)




				@slow
				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.tokenizer_class.from_pretrained("""bert-base-uncased"""		)

									_lowerCAmelCase						=    tokenizer.encode("""sequence builders"""						,				add_special_tokens=_snake_case		)
									_lowerCAmelCase						=    tokenizer.encode("""multi-sequence build"""						,				add_special_tokens=_snake_case		)

									_lowerCAmelCase						=    tokenizer.build_inputs_with_special_tokens(_snake_case		)
									_lowerCAmelCase						=    tokenizer.build_inputs_with_special_tokens(_snake_case						,				_snake_case		)

									assert encoded_sentence == [101] + text + [102]
									assert encoded_pair == [101] + text + [102] + text_a + [102]




				def 							snake_case   (							self		):



									"""simple docstring"""





									for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
														with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'		):
																			_lowerCAmelCase						=    self.rust_tokenizer_class.from_pretrained(_snake_case						,				**_snake_case		)

																			_lowerCAmelCase						=    F'A, naïve {tokenizer_r.mask_token} AllenNLP sentence.'
																			_lowerCAmelCase						=    tokenizer_r.encode_plus(
																			    _snake_case						,				return_attention_mask=_snake_case						,				return_token_type_ids=_snake_case						,				return_offsets_mapping=_snake_case						,				add_special_tokens=_snake_case						,				)

																			_lowerCAmelCase						=    tokenizer_r.do_lower_case if hasattr(_snake_case						,				"""do_lower_case"""		) else False
																			_lowerCAmelCase						=    (
																			    [
																			        ((0, 0), tokenizer_r.cls_token),
																			        ((0, 1), """A"""),
																			        ((1, 2), ""","""),
																			        ((3, 5), """na"""),
																			        ((5, 6), """##ï"""),
																			        ((6, 8), """##ve"""),
																			        ((9, 15), tokenizer_r.mask_token),
																			        ((16, 21), """Allen"""),
																			        ((21, 23), """##NL"""),
																			        ((23, 24), """##P"""),
																			        ((25, 33), """sentence"""),
																			        ((33, 34), """."""),
																			        ((0, 0), tokenizer_r.sep_token),
																			    ]
																			    if not do_lower_case
																			    else [
																			        ((0, 0), tokenizer_r.cls_token),
																			        ((0, 1), """a"""),
																			        ((1, 2), ""","""),
																			        ((3, 8), """naive"""),
																			        ((9, 15), tokenizer_r.mask_token),
																			        ((16, 21), """allen"""),
																			        ((21, 23), """##nl"""),
																			        ((23, 24), """##p"""),
																			        ((25, 33), """sentence"""),
																			        ((33, 34), """."""),
																			        ((0, 0), tokenizer_r.sep_token),
																			    ]
																			)

																			self.assertEqual(
																			    [e[1] for e in expected_results]						,				tokenizer_r.convert_ids_to_tokens(tokens["""input_ids"""]		)		)
																			self.assertEqual([e[0] for e in expected_results]						,				tokens["""offset_mapping"""]		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    ["""的""", """人""", """有"""]
									_lowerCAmelCase						=    """""".join(_snake_case		)
									for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
														with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'		):
																			_lowerCAmelCase						=    True
																			_lowerCAmelCase						=    self.tokenizer_class.from_pretrained(_snake_case						,				**_snake_case		)
																			_lowerCAmelCase						=    self.rust_tokenizer_class.from_pretrained(_snake_case						,				**_snake_case		)

																			_lowerCAmelCase						=    tokenizer_p.encode(_snake_case						,				add_special_tokens=_snake_case		)
																			_lowerCAmelCase						=    tokenizer_r.encode(_snake_case						,				add_special_tokens=_snake_case		)

																			_lowerCAmelCase						=    tokenizer_r.convert_ids_to_tokens(_snake_case		)
																			_lowerCAmelCase						=    tokenizer_p.convert_ids_to_tokens(_snake_case		)

																			# it is expected that each Chinese character is not preceded by "##"
																			self.assertListEqual(_snake_case						,				_snake_case		)
																			self.assertListEqual(_snake_case						,				_snake_case		)

																			_lowerCAmelCase						=    False
																			_lowerCAmelCase						=    self.rust_tokenizer_class.from_pretrained(_snake_case						,				**_snake_case		)
																			_lowerCAmelCase						=    self.tokenizer_class.from_pretrained(_snake_case						,				**_snake_case		)

																			_lowerCAmelCase						=    tokenizer_r.encode(_snake_case						,				add_special_tokens=_snake_case		)
																			_lowerCAmelCase						=    tokenizer_p.encode(_snake_case						,				add_special_tokens=_snake_case		)

																			_lowerCAmelCase						=    tokenizer_r.convert_ids_to_tokens(_snake_case		)
																			_lowerCAmelCase						=    tokenizer_p.convert_ids_to_tokens(_snake_case		)

																			# it is expected that only the first Chinese character is not preceded by "##".
																			_lowerCAmelCase						=    [
																			    F'##{token}' if idx != 0 else token for idx, token in enumerate(_snake_case		)
																			]
																			self.assertListEqual(_snake_case						,				_snake_case		)
																			self.assertListEqual(_snake_case						,				_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'''
   )
   A__:    Dict					    =       None
A__:    Tuple					    =       {
    '''7B''': 1_1008,
    '''13B''': 1_3824,
    '''30B''': 1_7920,
    '''65B''': 2_2016,
    '''70B''': 2_8672,
}
A__:    Any					    =       {
    '''7B''': 1,
    '''7Bf''': 1,
    '''13B''': 2,
    '''13Bf''': 2,
    '''30B''': 4,
    '''65B''': 8,
    '''70B''': 8,
    '''70Bf''': 8,
}


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[int]    ,_UpperCAmelCase	:	Optional[int]=1    ,_UpperCAmelCase	:	List[str]=256  )  ->			Dict:
       return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3  )  ) + multiple_of - 1) // multiple_of)


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	List[Any]  )  ->			List[str]:
       with open(_UpperCAmelCase    ,"""r"""  ) as f:
              return json.load(_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	str    ,_UpperCAmelCase	:	Optional[Any]  )  ->			Tuple:
       with open(_UpperCAmelCase    ,"""w"""  ) as f:
              json.dump(_UpperCAmelCase    ,_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	str    ,_UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	int    ,_UpperCAmelCase	:	List[Any]=True  )  ->			Union[str, Any]:
       os.makedirs(_UpperCAmelCase    ,exist_ok=_UpperCAmelCase  )
       _a    : Union[str, Any]     					=os.path.join(_UpperCAmelCase    ,"""tmp"""  )
       os.makedirs(_UpperCAmelCase    ,exist_ok=_UpperCAmelCase  )

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

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

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

                     _a    : Tuple     					={
                         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(),
                     }
                     _a    : str     					=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  )  )
                     _a    : Tuple     					=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    ,)
                     _a    : Any     					=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  )

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

              _a    : str     					=inv_freq
              for k, v in state_dict.items():
                     _a    : Any     					=filename
                     param_count += v.numel()
              torch.save(_UpperCAmelCase    ,os.path.join(_UpperCAmelCase    ,_UpperCAmelCase  )  )

       _a    : Union[str, Any]     					=F"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
       if model_size == "7B":
              # Unsharded
              _a    : List[str]     					={
                  """model.embed_tokens.weight""": loaded["""tok_embeddings.weight"""],
                  """model.norm.weight""": loaded["""norm.weight"""],
                  """lm_head.weight""": loaded["""output.weight"""],
              }
       else:
              _a    : 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():
              _a    : Dict     					=filename
              param_count += v.numel()
       torch.save(_UpperCAmelCase    ,os.path.join(_UpperCAmelCase    ,_UpperCAmelCase  )  )

       # Write configs
       _a    : Tuple     					={"""total_size""": param_count * 2}
       write_json(_UpperCAmelCase    ,os.path.join(_UpperCAmelCase    ,"""pytorch_model.bin.index.json"""  )  )
       _a    : Optional[Any]     					=params["""ffn_dim_multiplier"""] if """ffn_dim_multiplier""" in params else 1
       _a    : int     					=params["""multiple_of"""] if """multiple_of""" in params else 256
       _a    : List[Any]     					=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."""  )
       _a    : Any     					=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 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	int    ,_UpperCAmelCase	:	int  )  ->			Optional[Any]:
       # Initialize the tokenizer based on the `spm` model
       _a    : List[str]     					=LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
       print(F"Saving a {tokenizer_class.__name__} to {tokenizer_path}."  )
       _a    : List[Any]     					=tokenizer_class(_UpperCAmelCase  )
       tokenizer.save_pretrained(_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  )  ->			Union[str, Any]:
       _a    : 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`."""  )
       _a    : Optional[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    ,)
       _a    : List[Any]     					=os.path.join(args.input_dir    ,"""tokenizer.model"""  )
       write_tokenizer(args.output_dir    ,_UpperCAmelCase  )


if __name__ == "__main__":
   main()


