






"""simple docstring"""







import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SPIECE_UNDERLINE, is_sentencepiece_available
from transformers.models.speech_to_text import SpeechaTextTokenizer
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_UpperCamelCase					      =  get_tests_dir("""fixtures/test_sentencepiece.model""")

if is_sentencepiece_available():
      import sentencepiece as sp


_UpperCamelCase					      =  5
_UpperCamelCase					      =  10







@require_sentencepiece
@require_tokenizers
class  lowerCamelCase__ (	snake_case   ,					unittest.TestCase			):
    SCREAMING_SNAKE_CASE		=      SpeechaTextTokenizer
    SCREAMING_SNAKE_CASE		=      False
    SCREAMING_SNAKE_CASE		=      True






    def     _UpperCamelCase							(			self ):
          super().setUp()

          UpperCAmelCase          =	sp.SentencePieceProcessor()
          spm_model.Load(A )
          UpperCAmelCase          =	["""<s>""", """<pad>""", """</s>""", """<unk>"""]

          vocab += [spm_model.IdToPiece(id_ ) for id_ in range(len(A ) )]
          UpperCAmelCase          =	dict(zip(A     ,range(len(A ) ) ) )

          UpperCAmelCase          =	Path(self.tmpdirname )
          save_json(A     ,save_dir / VOCAB_FILES_NAMES["""vocab_file"""] )
          if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
                copyfile(A     ,save_dir / VOCAB_FILES_NAMES["""spm_file"""] )

          UpperCAmelCase          =	SpeechaTextTokenizer.from_pretrained(self.tmpdirname )
          tokenizer.save_pretrained(self.tmpdirname )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	"""<pad>"""
          UpperCAmelCase          =	1

          self.assertEqual(self.get_tokenizer()._convert_token_to_id(A )     ,A )
          self.assertEqual(self.get_tokenizer()._convert_id_to_token(A )     ,A )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	list(self.get_tokenizer().get_vocab().keys() )

          self.assertEqual(vocab_keys[0]     ,"""<s>""" )
          self.assertEqual(vocab_keys[1]     ,"""<pad>""" )
          self.assertEqual(vocab_keys[-1]     ,"""j""" )
          self.assertEqual(len(A )     ,1_001 )






    def     _UpperCamelCase							(			self ):
          self.assertEqual(self.get_tokenizer().vocab_size     ,1_001 )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	SpeechaTextTokenizer.from_pretrained(self.tmpdirname )

          UpperCAmelCase          =	tokenizer.tokenize("""This is a test""" )
          self.assertListEqual(A     ,["""▁This""", """▁is""", """▁a""", """▁t""", """est"""] )

          self.assertListEqual(
              tokenizer.convert_tokens_to_ids(A )     ,[289, 50, 14, 174, 386]     ,)

          UpperCAmelCase          =	tokenizer.tokenize("""I was born in 92000, and this is falsé.""" )
          self.assertListEqual(
              A     ,[SPIECE_UNDERLINE + """I""", SPIECE_UNDERLINE + """was""", SPIECE_UNDERLINE + """b""", """or""", """n""", SPIECE_UNDERLINE + """in""", SPIECE_UNDERLINE + """""", """9""", """2""", """0""", """0""", """0""", """,""", SPIECE_UNDERLINE + """and""", SPIECE_UNDERLINE + """this""", SPIECE_UNDERLINE + """is""", SPIECE_UNDERLINE + """f""", """al""", """s""", """é""", """."""]     ,)
          UpperCAmelCase          =	tokenizer.convert_tokens_to_ids(A )
          self.assertListEqual(A     ,[12, 25, 88, 59, 28, 23, 11, 4, 606, 351, 351, 351, 7, 16, 70, 50, 76, 84, 10, 4, 8] )

          UpperCAmelCase          =	tokenizer.convert_ids_to_tokens(A )
          self.assertListEqual(
              A     ,[SPIECE_UNDERLINE + """I""", SPIECE_UNDERLINE + """was""", SPIECE_UNDERLINE + """b""", """or""", """n""", SPIECE_UNDERLINE + """in""", SPIECE_UNDERLINE + """""", """<unk>""", """2""", """0""", """0""", """0""", """,""", SPIECE_UNDERLINE + """and""", SPIECE_UNDERLINE + """this""", SPIECE_UNDERLINE + """is""", SPIECE_UNDERLINE + """f""", """al""", """s""", """<unk>""", """."""]     ,)






    @slow
    def     _UpperCamelCase							(			self ):
          # fmt: off
          UpperCAmelCase          =	{"""input_ids""": [[3_791, 797, 31, 11, 64, 797, 31, 2_429, 433, 12, 1_176, 12, 20, 786, 915, 142, 2_413, 240, 37, 3_238, 797, 31, 11, 35, 93, 915, 142, 2_413, 240, 37, 5_540, 567, 1_276, 93, 37, 610, 40, 62, 455, 657, 1_042, 123, 780, 177, 37, 309, 241, 1_298, 514, 20, 292, 2_737, 114, 2_469, 241, 85, 64, 302, 548, 528, 423, 4, 509, 406, 423, 37, 601, 4, 777, 302, 548, 528, 423, 284, 4, 3_388, 511, 459, 4, 3_555, 40, 321, 302, 705, 4, 3_388, 511, 583, 326, 5, 5, 5, 62, 3_310, 560, 177, 2_680, 217, 1_508, 32, 31, 853, 418, 64, 583, 511, 1_605, 62, 35, 93, 560, 177, 2_680, 217, 1_508, 1_521, 64, 583, 511, 519, 62, 20, 1_515, 764, 20, 149, 261, 5_625, 7_972, 20, 5_540, 567, 1_276, 93, 3_925, 1_675, 11, 15, 802, 7_972, 576, 217, 1_508, 11, 35, 93, 1_253, 2_441, 15, 289, 652, 31, 416, 321, 3_842, 115, 40, 911, 8, 476, 619, 4, 380, 142, 423, 335, 240, 35, 93, 264, 8, 11, 335, 569, 420, 163, 5, 2], [260, 548, 528, 423, 20, 451, 20, 2_681, 1_153, 3_434, 20, 5_540, 37, 567, 126, 1_253, 2_441, 3_376, 449, 210, 431, 1_563, 177, 767, 5_540, 11, 1_203, 472, 11, 2_953, 685, 285, 364, 706, 1_153, 20, 6_799, 20, 2_869, 20, 4_464, 126, 40, 2_429, 20, 1_040, 866, 2_664, 418, 20, 318, 20, 1_726, 186, 20, 265, 522, 35, 93, 2_191, 4_634, 20, 1_040, 12, 6_799, 15, 228, 2_356, 142, 31, 11, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [2_575, 2_666, 684, 1_582, 1_176, 12, 627, 149, 619, 20, 4_902, 563, 11, 20, 149, 261, 3_420, 2_356, 174, 142, 4_714, 131, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], """attention_mask""": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
          # fmt: on

          self.tokenizer_integration_test_util(
              expected_encoding=A     ,model_name="""facebook/s2t-small-mustc-en-de-st"""     ,revision="""a14f04cf0776c02f62a8cb800cf7909e15ea23ad"""     ,)







@require_sentencepiece
class  lowerCamelCase__ (	unittest.TestCase			):
    SCREAMING_SNAKE_CASE		=      '''valhalla/s2t_mustc_multilinguial_medium'''

    SCREAMING_SNAKE_CASE		=      '''C\'est trop cool'''
    SCREAMING_SNAKE_CASE		=      '''Esto es genial'''






    @classmethod
    def     _UpperCamelCase							(			cls ):
          UpperCAmelCase          =	SpeechaTextTokenizer.from_pretrained(cls.checkpoint_name )
          return cls






    def     _UpperCamelCase							(			self ):
          self.assertEqual(self.tokenizer.lang_code_to_id["""pt"""]     ,4 )
          self.assertEqual(self.tokenizer.lang_code_to_id["""ru"""]     ,6 )
          self.assertEqual(self.tokenizer.lang_code_to_id["""it"""]     ,9 )
          self.assertEqual(self.tokenizer.lang_code_to_id["""de"""]     ,11 )






    def     _UpperCamelCase							(			self ):
          self.assertEqual(self.tokenizer.vocab_size     ,10_000 )






    def     _UpperCamelCase							(			self ):
          self.assertIn(A     ,self.tokenizer.all_special_ids )
          UpperCAmelCase          =	[ES_CODE, 4, 1_601, 47, 7_647, 2]
          UpperCAmelCase          =	self.tokenizer.decode(A     ,skip_special_tokens=A )
          UpperCAmelCase          =	self.tokenizer.decode(generated_ids[1:]     ,skip_special_tokens=A )
          self.assertEqual(A     ,A )
          self.assertNotIn(self.tokenizer.eos_token     ,A )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	"""fr"""
          UpperCAmelCase          =	self.tokenizer(self.french_text ).input_ids
          self.assertEqual(encoded[0]     ,A )
          self.assertEqual(encoded[-1]     ,self.tokenizer.eos_token_id )






    def     _UpperCamelCase							(			self ):
          UpperCAmelCase          =	"""fr"""
          self.assertListEqual(self.tokenizer.prefix_tokens     ,[FR_CODE] )

          UpperCAmelCase          =	"""es"""
          self.assertListEqual(self.tokenizer.prefix_tokens     ,[ES_CODE] )











"""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					      =  None
_UpperCamelCase					      =  {
    """7B""": 11008,
    """13B""": 13824,
    """30B""": 17920,
    """65B""": 22016,
    """70B""": 28672,
}
_UpperCamelCase					      =  {
    """7B""": 1,
    """7Bf""": 1,
    """13B""": 2,
    """13Bf""": 2,
    """30B""": 4,
    """65B""": 8,
    """70B""": 8,
    """70Bf""": 8,
}

def _a			(		_snake_case				,	_snake_case=1				,	_snake_case=256							):







      """simple docstring"""






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

def _a			(		_snake_case							):







      """simple docstring"""






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

def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






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

def _a			(		_snake_case				,	_snake_case				,	_snake_case				,	_snake_case=True							):







      """simple docstring"""






      os.makedirs(_snake_case				,	exist_ok=_snake_case							)
      UpperCAmelCase          =	os.path.join(_snake_case				,	"""tmp"""							)
      os.makedirs(_snake_case				,	exist_ok=_snake_case							)

      UpperCAmelCase          =	read_json(os.path.join(_snake_case				,	"""params.json"""							)							)
      UpperCAmelCase          =	NUM_SHARDS[model_size]
      UpperCAmelCase          =	params["""n_layers"""]
      UpperCAmelCase          =	params["""n_heads"""]
      UpperCAmelCase          =	n_heads // num_shards
      UpperCAmelCase          =	params["""dim"""]
      UpperCAmelCase          =	dim // n_heads
      UpperCAmelCase          =	10000.0
      UpperCAmelCase          =	1.0 / (base ** (torch.arange(0				,	_snake_case				,	2							).float() / dims_per_head))

      if "n_kv_heads" in params:
            UpperCAmelCase          =	params["""n_kv_heads"""]  # for GQA / MQA
            UpperCAmelCase          =	n_heads_per_shard // num_key_value_heads
            UpperCAmelCase          =	dim // num_key_value_heads
      else:  # compatibility with other checkpoints
            UpperCAmelCase          =	n_heads
            UpperCAmelCase          =	n_heads_per_shard
            UpperCAmelCase          =	dim

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

      print(F'''Fetching all parameters from the checkpoint at {input_base_path}.'''							)
      # Load weights
      if model_size == "7B":
            # Not sharded
            # (The sharded implementation would also work, but this is simpler.)
            UpperCAmelCase          =	torch.load(os.path.join(_snake_case				,	"""consolidated.00.pth"""							)				,	map_location="""cpu"""							)
      else:
            # Sharded
            UpperCAmelCase          =	[
                torch.load(os.path.join(_snake_case				,	F'''consolidated.{i:02d}.pth'''							)				,	map_location="""cpu"""							)
                for i in range(_snake_case							)
            ]
      UpperCAmelCase          =	0
      UpperCAmelCase          =	{"""weight_map""": {}}
      for layer_i in range(_snake_case							):
            UpperCAmelCase          =	F'''pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin'''
            if model_size == "7B":
                  # Unsharded
                  UpperCAmelCase          =	{
                      F'''model.layers.{layer_i}.self_attn.q_proj.weight''': permute(
                          loaded[F'''layers.{layer_i}.attention.wq.weight''']							),
                      F'''model.layers.{layer_i}.self_attn.k_proj.weight''': permute(
                          loaded[F'''layers.{layer_i}.attention.wk.weight''']							),
                      F'''model.layers.{layer_i}.self_attn.v_proj.weight''': loaded[F'''layers.{layer_i}.attention.wv.weight'''],
                      F'''model.layers.{layer_i}.self_attn.o_proj.weight''': loaded[F'''layers.{layer_i}.attention.wo.weight'''],
                      F'''model.layers.{layer_i}.mlp.gate_proj.weight''': loaded[F'''layers.{layer_i}.feed_forward.w1.weight'''],
                      F'''model.layers.{layer_i}.mlp.down_proj.weight''': loaded[F'''layers.{layer_i}.feed_forward.w2.weight'''],
                      F'''model.layers.{layer_i}.mlp.up_proj.weight''': loaded[F'''layers.{layer_i}.feed_forward.w3.weight'''],
                      F'''model.layers.{layer_i}.input_layernorm.weight''': loaded[F'''layers.{layer_i}.attention_norm.weight'''],
                      F'''model.layers.{layer_i}.post_attention_layernorm.weight''': loaded[F'''layers.{layer_i}.ffn_norm.weight'''],
                  }
            else:
                  # Sharded
                  # Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share
                  # the same storage object, saving attention_norm and ffn_norm will save other weights too, which is
                  # redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned.

                  UpperCAmelCase          =	{
                      F'''model.layers.{layer_i}.input_layernorm.weight''': loaded[0][
                          F'''layers.{layer_i}.attention_norm.weight'''
                      ].clone(),
                      F'''model.layers.{layer_i}.post_attention_layernorm.weight''': loaded[0][
                          F'''layers.{layer_i}.ffn_norm.weight'''
                      ].clone(),
                  }
                  UpperCAmelCase          =	permute(
                      torch.cat(
                          [
                              loaded[i][F'''layers.{layer_i}.attention.wq.weight'''].view(_snake_case				,	_snake_case				,	_snake_case							)
                              for i in range(_snake_case							)
                          ]				,	dim=0				,	).reshape(_snake_case				,	_snake_case							)							)
                  UpperCAmelCase          =	permute(
                      torch.cat(
                          [
                              loaded[i][F'''layers.{layer_i}.attention.wk.weight'''].view(
                                  _snake_case				,	_snake_case				,	_snake_case							)
                              for i in range(_snake_case							)
                          ]				,	dim=0				,	).reshape(_snake_case				,	_snake_case							)				,	_snake_case				,	_snake_case				,	_snake_case				,	)
                  UpperCAmelCase          =	torch.cat(
                      [
                          loaded[i][F'''layers.{layer_i}.attention.wv.weight'''].view(
                              _snake_case				,	_snake_case				,	_snake_case							)
                          for i in range(_snake_case							)
                      ]				,	dim=0				,	).reshape(_snake_case				,	_snake_case							)

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

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

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

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

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

      # Make space so we can load the model properly now.
      del state_dict
      del loaded
      gc.collect()

      print("""Loading the checkpoint in a Llama model."""							)
      UpperCAmelCase          =	LlamaForCausalLM.from_pretrained(_snake_case				,	torch_dtype=torch.floataa				,	low_cpu_mem_usage=_snake_case							)
      # Avoid saving this as part of the config.
      del model.config._name_or_path

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

def _a			(		_snake_case				,	_snake_case							):







      """simple docstring"""






      UpperCAmelCase          =	LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
      print(F'''Saving a {tokenizer_class.__name__} to {tokenizer_path}.'''							)
      UpperCAmelCase          =	tokenizer_class(_snake_case							)
      tokenizer.save_pretrained(_snake_case							)

def _a			(		):







      """simple docstring"""






      UpperCAmelCase          =	argparse.ArgumentParser()
      parser.add_argument(
          """--input_dir"""				,	help="""Location of LLaMA weights, which contains tokenizer.model and model folders"""				,	)
      parser.add_argument(
          """--model_size"""				,	choices=["""7B""", """7Bf""", """13B""", """13Bf""", """30B""", """65B""", """70B""", """70Bf""", """tokenizer_only"""]				,	)
      parser.add_argument(
          """--output_dir"""				,	help="""Location to write HF model and tokenizer"""				,	)
      parser.add_argument("""--safe_serialization"""				,	type=_snake_case				,	help="""Whether or not to save using `safetensors`."""							)
      UpperCAmelCase          =	parser.parse_args()
      if args.model_size != "tokenizer_only":
            write_model(
                model_path=args.output_dir				,	input_base_path=os.path.join(args.input_dir				,	args.model_size							)				,	model_size=args.model_size				,	safe_serialization=args.safe_serialization				,	)
      UpperCAmelCase          =	os.path.join(args.input_dir				,	"""tokenizer.model"""							)
      write_tokenizer(args.output_dir				,	_snake_case							)


if __name__ == "__main__":
  main()



