





import collections
from typing import List, Optional, Union

from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
from ..bert.tokenization_bert_fast import BertTokenizerFast
from .tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer, DPRReaderTokenizer


lowerCAmelCase_    =    logging.get_logger(__name__)

lowerCAmelCase_    =    {'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

lowerCAmelCase_    =    {
    '''vocab_file''': {
        '''facebook/dpr-ctx_encoder-single-nq-base''': (
            '''https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/vocab.txt'''
        ),
        '''facebook/dpr-ctx_encoder-multiset-base''': (
            '''https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/vocab.txt'''
        ),
    },
    '''tokenizer_file''': {
        '''facebook/dpr-ctx_encoder-single-nq-base''': (
            '''https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/tokenizer.json'''
        ),
        '''facebook/dpr-ctx_encoder-multiset-base''': (
            '''https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/tokenizer.json'''
        ),
    },
}
lowerCAmelCase_    =    {
    '''vocab_file''': {
        '''facebook/dpr-question_encoder-single-nq-base''': (
            '''https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/vocab.txt'''
        ),
        '''facebook/dpr-question_encoder-multiset-base''': (
            '''https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/vocab.txt'''
        ),
    },
    '''tokenizer_file''': {
        '''facebook/dpr-question_encoder-single-nq-base''': (
            '''https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/tokenizer.json'''
        ),
        '''facebook/dpr-question_encoder-multiset-base''': (
            '''https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/tokenizer.json'''
        ),
    },
}
lowerCAmelCase_    =    {
    '''vocab_file''': {
        '''facebook/dpr-reader-single-nq-base''': (
            '''https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/vocab.txt'''
        ),
        '''facebook/dpr-reader-multiset-base''': (
            '''https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/vocab.txt'''
        ),
    },
    '''tokenizer_file''': {
        '''facebook/dpr-reader-single-nq-base''': (
            '''https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/tokenizer.json'''
        ),
        '''facebook/dpr-reader-multiset-base''': (
            '''https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/tokenizer.json'''
        ),
    },
}

lowerCAmelCase_    =    {
    '''facebook/dpr-ctx_encoder-single-nq-base''': 5_1_2,
    '''facebook/dpr-ctx_encoder-multiset-base''': 5_1_2,
}
lowerCAmelCase_    =    {
    '''facebook/dpr-question_encoder-single-nq-base''': 5_1_2,
    '''facebook/dpr-question_encoder-multiset-base''': 5_1_2,
}
lowerCAmelCase_    =    {
    '''facebook/dpr-reader-single-nq-base''': 5_1_2,
    '''facebook/dpr-reader-multiset-base''': 5_1_2,
}


lowerCAmelCase_    =    {
    '''facebook/dpr-ctx_encoder-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-ctx_encoder-multiset-base''': {'''do_lower_case''': True},
}
lowerCAmelCase_    =    {
    '''facebook/dpr-question_encoder-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-question_encoder-multiset-base''': {'''do_lower_case''': True},
}
lowerCAmelCase_    =    {
    '''facebook/dpr-reader-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-reader-multiset-base''': {'''do_lower_case''': True},
}



class  __lowerCAmelCase	(      a_					):
			lowerCamelCase_					:					Any	      =				VOCAB_FILES_NAMES
			lowerCamelCase_					:					List[Any]	      =				CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
			lowerCamelCase_					:					Any	      =				CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCamelCase_					:					List[Any]	      =				CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION
			lowerCamelCase_					:					Optional[Any]	      =				DPRContextEncoderTokenizer




class  __lowerCAmelCase	(      a_					):
			lowerCamelCase_					:					int	      =				VOCAB_FILES_NAMES
			lowerCamelCase_					:					List[Any]	      =				QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
			lowerCamelCase_					:					Union[str, Any]	      =				QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCamelCase_					:					Tuple	      =				QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION
			lowerCamelCase_					:					List[str]	      =				DPRQuestionEncoderTokenizer




lowerCAmelCase_    =    collections.namedtuple(
    '''DPRSpanPrediction''', ['''span_score''', '''relevance_score''', '''doc_id''', '''start_index''', '''end_index''', '''text''']
)

lowerCAmelCase_    =    collections.namedtuple('''DPRReaderOutput''', ['''start_logits''', '''end_logits''', '''relevance_logits'''])


lowerCAmelCase_    =    r'''
    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.
    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),
    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`
    with the format:

    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>

    Args:
        questions (`str` or `List[str]`):
            The questions to be encoded. You can specify one question for many passages. In this case, the question
            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in
            `titles` or `texts`.
        titles (`str` or `List[str]`):
            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.
        texts (`str` or `List[str]`):
            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
            Activates and controls padding. Accepts the following values:

            - `True` or `\'longest\'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
              if provided).
            - `\'max_length\'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `\'do_not_pad\'` (default): No padding (i.e., can output a batch with sequences of different
              lengths).
        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):
            Activates and controls truncation. Accepts the following values:

            - `True` or `\'longest_first\'`: Truncate to a maximum length specified with the argument `max_length` or to
              the maximum acceptable input length for the model if that argument is not provided. This will truncate
              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch
              of pairs) is provided.
            - `\'only_first\'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the first
              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `\'only_second\'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the
              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `False` or `\'do_not_truncate\'` (default): No truncation (i.e., can output batch with sequence lengths
              greater than the model maximum admissible input size).
        max_length (`int`, *optional*):
                Controls the maximum length to use by one of the truncation/padding parameters.

                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
                is required by one of the truncation/padding parameters. If the model has no specific maximum input
                length (like XLNet) truncation/padding to a maximum length will be deactivated.
        return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors instead of list of python integers. Acceptable values are:

                - `\'tf\'`: Return TensorFlow `tf.constant` objects.
                - `\'pt\'`: Return PyTorch `torch.Tensor` objects.
                - `\'np\'`: Return Numpy `np.ndarray` objects.
        return_attention_mask (`bool`, *optional*):
            Whether or not to return the attention mask. If not set, will return the attention mask according to the
            specific tokenizer\'s default, defined by the `return_outputs` attribute.

            [What are attention masks?](../glossary#attention-mask)

    Return:
        `Dict[str, List[List[int]]]`: A dictionary with the following keys:

        - `input_ids`: List of token ids to be fed to a model.
        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.
    '''



@add_start_docstrings(a_					)
class  __lowerCAmelCase	:






			def __call__(self ,							__magic_name__ ,							__magic_name__ = None ,							__magic_name__ = None ,							__magic_name__ = False ,							__magic_name__ = False ,							__magic_name__ = None ,							__magic_name__ = None ,							__magic_name__ = None ,							**__magic_name__ ,							)					->       BatchEncoding:


							'''simple docstring'''



							if titles is None and texts is None:
											return super().__call__(
											    lowercase_ ,							padding=lowercase_ ,							truncation=lowercase_ ,							max_length=lowercase_ ,							return_tensors=lowercase_ ,							return_attention_mask=lowercase_ ,							**lowercase_ ,							)
							elif titles is None or texts is None:
											snake_case_						:   Any           =  titles if texts is None else texts
											return super().__call__(
											    lowercase_ ,							lowercase_ ,							padding=lowercase_ ,							truncation=lowercase_ ,							max_length=lowercase_ ,							return_tensors=lowercase_ ,							return_attention_mask=lowercase_ ,							**lowercase_ ,							)
							snake_case_						:   int           =  titles if not isinstance(lowercase_ ,							lowercase_   ) else [titles]
							snake_case_						:   Any           =  texts if not isinstance(lowercase_ ,							lowercase_   ) else [texts]
							snake_case_						:   str           =  len(lowercase_   )
							snake_case_						:   Any           =  questions if not isinstance(lowercase_ ,							lowercase_   ) else [questions] * n_passages
							assert len(lowercase_   ) == len(
							    lowercase_   ), F'''There should be as many titles than texts but got {len(lowercase_   )} titles and {len(lowercase_   )} texts.'''
							snake_case_						:   str           =  super().__call__(lowercase_ ,							lowercase_ ,							padding=lowercase_ ,							truncation=lowercase_   )['''input_ids''']
							snake_case_						:   List[Any]           =  super().__call__(lowercase_ ,							add_special_tokens=lowercase_ ,							padding=lowercase_ ,							truncation=lowercase_   )['''input_ids''']
							snake_case_						:   List[Any]           =  {
							    '''input_ids''': [
							        (encoded_question_and_title + encoded_text)[:max_length]
							        if max_length is not None and truncation
							        else encoded_question_and_title + encoded_text
							        for encoded_question_and_title, encoded_text in zip(lowercase_ ,							lowercase_   )
							    ]
							}
							if return_attention_mask is not False:
											snake_case_						:   Optional[int]           =  []
											for input_ids in encoded_inputs["input_ids"]:
															attention_mask.append([int(input_id != self.pad_token_id   ) for input_id in input_ids]   )
											snake_case_						:   Dict           =  attention_mask
							return self.pad(lowercase_ ,							padding=lowercase_ ,							max_length=lowercase_ ,							return_tensors=lowercase_   )






			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ = 16 ,							__magic_name__ = 64 ,							__magic_name__ = 4 ,							)					->       List[DPRSpanPrediction]:


							'''simple docstring'''



							snake_case_						:   List[str]           =  reader_input['''input_ids''']
							snake_case_      ,							snake_case_      ,							snake_case_						:   Optional[Any]           =  reader_output[:3]
							snake_case_						:   Optional[Any]           =  len(lowercase_   )
							snake_case_						:   str           =  sorted(range(lowercase_   ) ,							reverse=lowercase_ ,							key=relevance_logits.__getitem__   )
							snake_case_						:   Optional[int]           =  []
							for doc_id in sorted_docs:
											snake_case_						:   List[str]           =  list(input_ids[doc_id]   )
											# assuming question & title information is at the beginning of the sequence
											snake_case_						:   str           =  sequence_ids.index(self.sep_token_id ,							2   ) + 1  # second sep id
											if sequence_ids[-1] == self.pad_token_id:
															snake_case_						:   int           =  sequence_ids.index(self.pad_token_id   )
											else:
															snake_case_						:   int           =  len(lowercase_   )

											snake_case_						:   List[Any]           =  self._get_best_spans(
											    start_logits=start_logits[doc_id][passage_offset:sequence_len] ,							end_logits=end_logits[doc_id][passage_offset:sequence_len] ,							max_answer_length=lowercase_ ,							top_spans=lowercase_ ,							)
											for start_index, end_index in best_spans:
															start_index += passage_offset
															end_index += passage_offset
															nbest_spans_predictions.append(
															    DPRSpanPrediction(
															        span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index] ,							relevance_score=relevance_logits[doc_id] ,							doc_id=lowercase_ ,							start_index=lowercase_ ,							end_index=lowercase_ ,							text=self.decode(sequence_ids[start_index : end_index + 1]   ) ,							)   )
											if len(lowercase_   ) >= num_spans:
															break
							return nbest_spans_predictions[:num_spans]






			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ ,							)					->       List[DPRSpanPrediction]:


							'''simple docstring'''



							snake_case_						:   Dict           =  []
							for start_index, start_score in enumerate(lowercase_   ):
											for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]   ):
															scores.append(((start_index, start_index + answer_length), start_score + end_score)   )
							snake_case_						:   List[Any]           =  sorted(lowercase_ ,							key=lambda __magic_name__   : x[1] ,							reverse=lowercase_   )
							snake_case_						:   int           =  []
							for (start_index, end_index), score in scores:
											assert start_index <= end_index, F'''Wrong span indices: [{start_index}:{end_index}]'''
											snake_case_						:   Dict           =  end_index - start_index + 1
											assert length <= max_answer_length, F'''Span is too long: {length} > {max_answer_length}'''
											if any(
											    start_index <= prev_start_index <= prev_end_index <= end_index
											    or prev_start_index <= start_index <= end_index <= prev_end_index
											    for (prev_start_index, prev_end_index) in chosen_span_intervals   ):
															continue
											chosen_span_intervals.append((start_index, end_index)   )

											if len(lowercase_   ) == top_spans:
															break
							return chosen_span_intervals




@add_end_docstrings(a_					)
class  __lowerCAmelCase	(      a_,					a_					):
			lowerCamelCase_					:					Any	      =				VOCAB_FILES_NAMES
			lowerCamelCase_					:					Dict	      =				READER_PRETRAINED_VOCAB_FILES_MAP
			lowerCamelCase_					:					Any	      =				READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCamelCase_					:					List[Any]	      =				READER_PRETRAINED_INIT_CONFIGURATION
			lowerCamelCase_					:					Tuple	      =				['''input_ids''', '''attention_mask''']
			lowerCamelCase_					:					Any	      =				DPRReaderTokenizer







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_1_0_0_8,
    '''13B''': 1_3_8_2_4,
    '''30B''': 1_7_9_2_0,
    '''65B''': 2_2_0_1_6,
    '''70B''': 2_8_6_7_2,
}
lowerCAmelCase_    =    {
    '''7B''': 1,
    '''7Bf''': 1,
    '''13B''': 2,
    '''13Bf''': 2,
    '''30B''': 4,
    '''65B''': 8,
    '''70B''': 8,
    '''70Bf''': 8,
}






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase=1			,    _UpperCamelCase=256			)		->						Optional[int]:






				"""simple docstring"""


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






def 					lowerCamelCase_						(    _UpperCamelCase			)		->						Optional[Any]:






				"""simple docstring"""


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






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase			)		->						int:






				"""simple docstring"""


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






def 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase			,    _UpperCamelCase			,    _UpperCamelCase=True			)		->						Optional[Any]:






				"""simple docstring"""


				os.makedirs(_UpperCamelCase			,    exist_ok=_UpperCamelCase			)
				snake_case_						:   int           =  os.path.join(_UpperCamelCase			,    '''tmp'''			)
				os.makedirs(_UpperCamelCase			,    exist_ok=_UpperCamelCase			)

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

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

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

												snake_case_						:   Union[str, Any]           =  {
												    f'''model.layers.{layer_i}.input_layernorm.weight''': loaded[0][
												        f'''layers.{layer_i}.attention_norm.weight'''
												    ].clone(),
												    f'''model.layers.{layer_i}.post_attention_layernorm.weight''': loaded[0][
												        f'''layers.{layer_i}.ffn_norm.weight'''
												    ].clone(),
												}
												snake_case_						:   int           =  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			)			)
												snake_case_						:   Optional[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			,    )
												snake_case_						:   int           =  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			)

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

								snake_case_						:   str           =  inv_freq
								for k, v in state_dict.items():
												snake_case_						:   Dict           =  filename
												param_count += v.numel()
								torch.save(_UpperCamelCase			,    os.path.join(_UpperCamelCase			,    _UpperCamelCase			)			)

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

				# Write configs
				snake_case_						:   int           =  {'''total_size''': param_count * 2}
				write_json(_UpperCamelCase			,    os.path.join(_UpperCamelCase			,    '''pytorch_model.bin.index.json'''			)			)
				snake_case_						:   str           =  params['''ffn_dim_multiplier'''] if '''ffn_dim_multiplier''' in params else 1
				snake_case_						:   Optional[int]           =  params['''multiple_of'''] if '''multiple_of''' in params else 256
				snake_case_						:   Optional[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.'''			)
				snake_case_						:   Union[str, 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 					lowerCamelCase_						(    _UpperCamelCase			,    _UpperCamelCase			)		->						List[str]:






				"""simple docstring"""


				snake_case_						:   Union[str, Any]           =  LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
				print(f'''Saving a {tokenizer_class.__name__} to {tokenizer_path}.'''			)
				snake_case_						:   Union[str, Any]           =  tokenizer_class(_UpperCamelCase			)
				tokenizer.save_pretrained(_UpperCamelCase			)






def 					lowerCamelCase_						(    )		->						int:






				"""simple docstring"""


				snake_case_						:   Union[str, Any]           =  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`.'''			)
				snake_case_						:   Dict           =  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			,    )
				snake_case_						:   Dict           =  os.path.join(args.input_dir			,    '''tokenizer.model'''			)
				write_tokenizer(args.output_dir			,    _UpperCamelCase			)


if __name__ == "__main__":
	main()
