






"""simple docstring"""







from __future__ import annotations

import copy
import inspect
import unittest

import numpy as np

from transformers import is_tf_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import tensorflow as tf

	from transformers import (
	    TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
	    TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
	    TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
	    TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
	    LayoutLMvaConfig,
	    TFLayoutLMvaForQuestionAnswering,
	    TFLayoutLMvaForSequenceClassification,
	    TFLayoutLMvaForTokenClassification,
	    TFLayoutLMvaModel,
	)

if is_vision_available():
	from PIL import Image

	from transformers import LayoutLMvaImageProcessor







class        _SCREAMING_SNAKE_CASE						:





							"""simple docstring"""





							def __init__(   self:     List[str]				,		__A:     str				,		__A:     Any=2				,		__A:     str=3				,		__A:     Tuple=4				,		__A:     Dict=2				,		__A:     List[Any]=7				,		__A:     Any=True				,		__A:     Any=True				,		__A:     List[str]=True				,		__A:     Optional[int]=True				,		__A:     Optional[int]=99				,		__A:     Tuple=36				,		__A:     List[str]=2				,		__A:     Dict=4				,		__A:     List[str]=37				,		__A:     Optional[int]="gelu"				,		__A:     Optional[int]=0.1				,		__A:     Tuple=0.1				,		__A:     List[Any]=512				,		__A:     List[str]=16				,		__A:     Any=2				,		__A:     Union[str, Any]=0.0_2				,		__A:     Optional[int]=6				,		__A:     Union[str, Any]=6				,		__A:     Union[str, Any]=3				,		__A:     Tuple=4				,		__A:     Optional[int]=None				,		__A:     Optional[Any]=1000				,		):

												'''simple docstring'''

												a__          =  parent
												a__          =  batch_size
												a__          =  num_channels
												a__          =  image_size
												a__          =  patch_size
												a__          =  is_training
												a__          =  use_input_mask
												a__          =  use_token_type_ids
												a__          =  use_labels
												a__          =  vocab_size
												a__          =  hidden_size
												a__          =  num_hidden_layers
												a__          =  num_attention_heads
												a__          =  intermediate_size
												a__          =  hidden_act
												a__          =  hidden_dropout_prob
												a__          =  attention_probs_dropout_prob
												a__          =  max_position_embeddings
												a__          =  type_vocab_size
												a__          =  type_sequence_label_size
												a__          =  initializer_range
												a__          =  coordinate_size
												a__          =  shape_size
												a__          =  num_labels
												a__          =  num_choices
												a__          =  scope
												a__          =  range_bbox

												# LayoutLMv3's sequence length equals the number of text tokens + number of patches + 1 (we add 1 for the CLS token)
												a__          =  text_seq_length
												a__          =  (image_size // patch_size) ** 2 + 1
												a__          =  self.text_seq_length + self.image_seq_length




							def  lowercase		(   self:     Optional[int]    ):

												'''simple docstring'''

												a__          =  ids_tensor([self.batch_size, self.text_seq_length]				,		self.vocab_size    )

												a__          =  ids_tensor([self.batch_size, self.text_seq_length, 4]				,		self.range_bbox    )
												a__          =  bbox.numpy()
												# Ensure that bbox is legal
												for i in range(bbox.shape[0]    ):
																	for j in range(bbox.shape[1]    ):
																						if bbox[i, j, 3] < bbox[i, j, 1]:
																											a__          =  bbox[i, j, 3]
																											a__          =  bbox[i, j, 1]
																											a__          =  tmp_coordinate
																						if bbox[i, j, 2] < bbox[i, j, 0]:
																											a__          =  bbox[i, j, 2]
																											a__          =  bbox[i, j, 0]
																											a__          =  tmp_coordinate
												a__          =  tf.constant(__A    )

												a__          =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

												a__          =  None
												if self.use_input_mask:
																	a__          =  random_attention_mask([self.batch_size, self.text_seq_length]    )

												a__          =  None
												if self.use_token_type_ids:
																	a__          =  ids_tensor([self.batch_size, self.text_seq_length]				,		self.type_vocab_size    )

												a__          =  None
												a__          =  None
												if self.use_labels:
																	a__          =  ids_tensor([self.batch_size]				,		self.type_sequence_label_size    )
																	a__          =  ids_tensor([self.batch_size, self.text_seq_length]				,		self.num_labels    )

												a__          =  LayoutLMvaConfig(
												    vocab_size=self.vocab_size				,		hidden_size=self.hidden_size				,		num_hidden_layers=self.num_hidden_layers				,		num_attention_heads=self.num_attention_heads				,		intermediate_size=self.intermediate_size				,		hidden_act=self.hidden_act				,		hidden_dropout_prob=self.hidden_dropout_prob				,		attention_probs_dropout_prob=self.attention_probs_dropout_prob				,		max_position_embeddings=self.max_position_embeddings				,		type_vocab_size=self.type_vocab_size				,		initializer_range=self.initializer_range				,		coordinate_size=self.coordinate_size				,		shape_size=self.shape_size				,		input_size=self.image_size				,		patch_size=self.patch_size				,		)

												return config, input_ids, bbox, pixel_values, token_type_ids, input_mask, sequence_labels, token_labels




							def  lowercase		(   self:     List[str]				,		__A:     List[str]				,		__A:     List[str]				,		__A:     List[str]				,		__A:     int				,		__A:     Any				,		__A:     Any    ):

												'''simple docstring'''

												a__          =  TFLayoutLMvaModel(config=__A    )

												# text + image
												a__          =  model(__A				,		pixel_values=__A				,		training=__A    )
												a__          =  model(
												    __A				,		bbox=__A				,		pixel_values=__A				,		attention_mask=__A				,		token_type_ids=__A				,		training=__A				,		)
												a__          =  model(__A				,		bbox=__A				,		pixel_values=__A				,		training=__A    )

												self.parent.assertEqual(result.last_hidden_state.shape				,		(self.batch_size, self.seq_length, self.hidden_size)    )

												# text only
												a__          =  model(__A				,		training=__A    )
												self.parent.assertEqual(
												    result.last_hidden_state.shape				,		(self.batch_size, self.text_seq_length, self.hidden_size)    )

												# image only
												a__          =  model({'''pixel_values''': pixel_values}				,		training=__A    )
												self.parent.assertEqual(
												    result.last_hidden_state.shape				,		(self.batch_size, self.image_seq_length, self.hidden_size)    )




							def  lowercase		(   self:     Optional[int]				,		__A:     Any				,		__A:     str				,		__A:     List[str]				,		__A:     List[str]				,		__A:     List[str]				,		__A:     Optional[Any]				,		__A:     Any    ):

												'''simple docstring'''

												a__          =  self.num_labels
												a__          =  TFLayoutLMvaForSequenceClassification(config=__A    )
												a__          =  model(
												    __A				,		bbox=__A				,		pixel_values=__A				,		attention_mask=__A				,		token_type_ids=__A				,		labels=__A				,		training=__A				,		)
												self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_labels)    )




							def  lowercase		(   self:     str				,		__A:     List[str]				,		__A:     int				,		__A:     str				,		__A:     Any				,		__A:     str				,		__A:     str				,		__A:     str    ):

												'''simple docstring'''

												a__          =  self.num_labels
												a__          =  TFLayoutLMvaForTokenClassification(config=__A    )
												a__          =  model(
												    __A				,		bbox=__A				,		pixel_values=__A				,		attention_mask=__A				,		token_type_ids=__A				,		labels=__A				,		training=__A				,		)
												self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.text_seq_length, self.num_labels)    )




							def  lowercase		(   self:     Union[str, Any]				,		__A:     Any				,		__A:     List[Any]				,		__A:     Any				,		__A:     List[str]				,		__A:     Any				,		__A:     Optional[int]				,		__A:     Tuple    ):

												'''simple docstring'''

												a__          =  2
												a__          =  TFLayoutLMvaForQuestionAnswering(config=__A    )
												a__          =  model(
												    __A				,		bbox=__A				,		pixel_values=__A				,		attention_mask=__A				,		token_type_ids=__A				,		start_positions=__A				,		end_positions=__A				,		training=__A				,		)
												self.parent.assertEqual(result.start_logits.shape				,		(self.batch_size, self.seq_length)    )
												self.parent.assertEqual(result.end_logits.shape				,		(self.batch_size, self.seq_length)    )





							def  lowercase		(   self:     int    ):

												'''simple docstring'''

												a__          =  self.prepare_config_and_inputs()
												((a__)       ,(a__)       ,(a__)       ,(a__)       ,(a__)       ,(a__)       ,(a__)       ,(a__))          =  config_and_inputs
												a__          =  {
												    '''input_ids''': input_ids,
												    '''bbox''': bbox,
												    '''pixel_values''': pixel_values,
												    '''token_type_ids''': token_type_ids,
												    '''attention_mask''': input_mask,
												}
												return config, inputs_dict







@require_tf
class        _SCREAMING_SNAKE_CASE						(  __snake_case       , __snake_case       , unittest.TestCase							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =(
							    (
							        TFLayoutLMvaModel,
							        TFLayoutLMvaForQuestionAnswering,
							        TFLayoutLMvaForSequenceClassification,
							        TFLayoutLMvaForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							_SCREAMING_SNAKE_CASE    =(
							    {'document-question-answering': TFLayoutLMvaForQuestionAnswering, 'feature-extraction': TFLayoutLMvaModel}
							    if is_tf_available()
							    else {}
							)

							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False




							def  lowercase		(   self:     List[str]				,		__A:     Dict				,		__A:     Optional[Any]				,		__A:     str				,		__A:     Union[str, Any]				,		__A:     Optional[Any]    ):

												'''simple docstring'''

												return True




							def  lowercase		(   self:     Dict				,		__A:     Optional[Any]				,		__A:     Any				,		__A:     List[Any]=False    ):

												'''simple docstring'''

												a__          =  copy.deepcopy(__A    )

												if model_class in get_values(__A    ):
																	a__          =  {
																	    k: tf.tile(tf.expand_dims(__A				,		1    )				,		(1, self.model_tester.num_choices) + (1,) * (v.ndim - 1)    )
																	    if isinstance(__A				,		tf.Tensor    ) and v.ndim > 0
																	    else v
																	    for k, v in inputs_dict.items()
																	}

												if return_labels:
																	if model_class in get_values(__A    ):
																						a__          =  tf.ones(self.model_tester.batch_size				,		dtype=tf.intaa    )
																	elif model_class in get_values(__A    ):
																						a__          =  tf.zeros(self.model_tester.batch_size				,		dtype=tf.intaa    )
																						a__          =  tf.zeros(self.model_tester.batch_size				,		dtype=tf.intaa    )
																	elif model_class in get_values(__A    ):
																						a__          =  tf.zeros(self.model_tester.batch_size				,		dtype=tf.intaa    )
																	elif model_class in get_values(__A    ):
																						a__          =  tf.zeros(
																						    (self.model_tester.batch_size, self.model_tester.text_seq_length)				,		dtype=tf.intaa    )

												return inputs_dict




							def  lowercase		(   self:     List[str]    ):

												'''simple docstring'''

												a__          =  TFLayoutLMvaModelTester(self    )
												a__          =  ConfigTester(self				,		config_class=__A				,		hidden_size=37    )




							def  lowercase		(   self:     Dict    ):

												'''simple docstring'''

												self.config_tester.run_common_tests()




							def  lowercase		(   self:     Any    ):

												'''simple docstring'''

												a__       ,a__          =  self.model_tester.prepare_config_and_inputs_for_common()
												for model_class in self.all_model_classes:
																	a__          =  model_class(__A    )
																	if getattr(__A				,		'''hf_compute_loss'''				,		__A    ):
																						# The number of elements in the loss should be the same as the number of elements in the label
																						a__          =  self._prepare_for_class(inputs_dict.copy()				,		__A				,		return_labels=__A    )
																						a__          =  prepared_for_class[
																						    sorted(prepared_for_class.keys() - inputs_dict.keys()				,		reverse=__A    )[0]
																						]
																						a__          =  added_label.shape.as_list()[:1]

																						# Test that model correctly compute the loss with kwargs
																						a__          =  self._prepare_for_class(inputs_dict.copy()				,		__A				,		return_labels=__A    )
																						a__          =  prepared_for_class.pop('''input_ids'''    )

																						a__          =  model(__A				,		**__A    )[0]
																						self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )

																						# Test that model correctly compute the loss when we mask some positions
																						a__          =  self._prepare_for_class(inputs_dict.copy()				,		__A				,		return_labels=__A    )
																						a__          =  prepared_for_class.pop('''input_ids'''    )
																						if "labels" in prepared_for_class:
																											a__          =  prepared_for_class['''labels'''].numpy()
																											if len(labels.shape    ) > 1 and labels.shape[1] != 1:
																																a__          =  -100
																																a__          =  tf.convert_to_tensor(__A    )
																																a__          =  model(__A				,		**__A    )[0]
																																self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )
																																self.assertTrue(not np.any(np.isnan(loss.numpy()    )    )    )

                # Test that model correctly compute the loss with a dict
																						a__          =  self._prepare_for_class(inputs_dict.copy()				,		__A				,		return_labels=__A    )
																						a__          =  model(__A    )[0]
																						self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )

																						# Test that model correctly compute the loss with a tuple
																						a__          =  self._prepare_for_class(inputs_dict.copy()				,		__A				,		return_labels=__A    )

																						# Get keys that were added with the _prepare_for_class function
																						a__          =  prepared_for_class.keys() - inputs_dict.keys()
																						a__          =  inspect.signature(model.call    ).parameters
																						a__          =  list(signature.keys()    )

																						# Create a dictionary holding the location of the tensors in the tuple
																						a__          =  {0: '''input_ids'''}
																						for label_key in label_keys:
																											a__          =  signature_names.index(__A    )
																											a__          =  label_key
																						a__          =  sorted(tuple_index_mapping.items()    )
																						# Initialize a list with their default values, update the values and convert to a tuple
																						a__          =  []

																						for name in signature_names:
																											if name != "kwargs":
																																list_input.append(signature[name].default    )

																						for index, value in sorted_tuple_index_mapping:
																											a__          =  prepared_for_class[value]

																						a__          =  tuple(__A    )

																						# Send to model
																						a__          =  model(tuple_input[:-1]    )[0]

																						self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]    )




							def  lowercase		(   self:     Optional[int]    ):

												'''simple docstring'''

												(
												    (
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,
												)          =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(__A				,		__A				,		__A				,		__A				,		__A				,		__A    )




							def  lowercase		(   self:     Dict    ):

												'''simple docstring'''

												(
												    (
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,
												)          =  self.model_tester.prepare_config_and_inputs()
												for type in ["absolute", "relative_key", "relative_key_query"]:
																	a__          =  type
																	self.model_tester.create_and_check_model(__A				,		__A				,		__A				,		__A				,		__A				,		__A    )




							def  lowercase		(   self:     int    ):

												'''simple docstring'''

												(
												    (
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,
												)          =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_sequence_classification(
												    __A				,		__A				,		__A				,		__A				,		__A				,		__A				,		__A    )




							def  lowercase		(   self:     List[str]    ):

												'''simple docstring'''

												(
												    (
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,
												)          =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_token_classification(
												    __A				,		__A				,		__A				,		__A				,		__A				,		__A				,		__A    )




							def  lowercase		(   self:     Tuple    ):

												'''simple docstring'''

												(
												    (
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,(
												    a__
												)       ,
												)          =  self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_question_answering(
												    __A				,		__A				,		__A				,		__A				,		__A				,		__A				,		__A    )





							@slow
							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												for model_name in TF_LAYOUTLMV3_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	a__          =  TFLayoutLMvaModel.from_pretrained(__A    )
																	self.assertIsNotNone(__A    )





def        SCREAMING_SNAKE_CASE       (      ):
					a__          =  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
					return image







@require_tf
class        _SCREAMING_SNAKE_CASE						(  unittest.TestCase							):





							"""simple docstring"""





							@cached_property
							def  lowercase		(   self:     Any    ):

												'''simple docstring'''

												return LayoutLMvaImageProcessor(apply_ocr=__A    ) if is_vision_available() else None





							@slow
							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  TFLayoutLMvaModel.from_pretrained('''microsoft/layoutlmv3-base'''    )

												a__          =  self.default_image_processor
												a__          =  prepare_img()
												a__          =  image_processor(images=__A				,		return_tensors='''tf'''    ).pixel_values

												a__          =  tf.constant([[1, 2]]    )
												a__          =  tf.expand_dims(tf.constant([[1, 2, 3, 4], [5, 6, 7, 8]]    )				,		axis=0    )

												# forward pass
												a__          =  model(input_ids=__A				,		bbox=__A				,		pixel_values=__A				,		training=__A    )

												# verify the logits
												a__          =  (1, 199, 768)
												self.assertEqual(outputs.last_hidden_state.shape				,		__A    )

												a__          =  tf.constant(
												    [[-0.0_5_2_9, 0.3_6_1_8, 0.1_6_3_2], [-0.1_5_8_7, -0.1_6_6_7, -0.0_4_0_0], [-0.1_5_5_7, -0.1_6_7_1, -0.0_5_0_5]]    )

												self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]				,		__A				,		atol=1e-4    )    )










"""simple docstring"""







import argparse
import gc
import json
import os
import shutil
import warnings

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
	from transformers import LlamaTokenizerFast
except ImportError as e:
	warnings.warn(e)
	warnings.warn(
	    'The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion'
	)
	__a       :	Union[str, Any]             = None
__a       :	Union[str, Any]             = {
    '7B': 1_1008,
    '13B': 1_3824,
    '30B': 1_7920,
    '65B': 2_2016,
    '70B': 2_8672,
}
__a       :	List[Any]             = {
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_=1 ,     lowerCamelCase_=256):
					return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of)





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_):
					with open(lowerCamelCase_ ,     '''r''') as f:
										return json.load(lowerCamelCase_)





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_):
					with open(lowerCamelCase_ ,     '''w''') as f:
										json.dump(lowerCamelCase_ ,     lowerCamelCase_)





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_ ,     lowerCamelCase_ ,     lowerCamelCase_=True):
					os.makedirs(lowerCamelCase_ ,     exist_ok=lowerCamelCase_)
					a__          =  os.path.join(lowerCamelCase_ ,     '''tmp''')
					os.makedirs(lowerCamelCase_ ,     exist_ok=lowerCamelCase_)

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

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

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

					print(f'Fetching all parameters from the checkpoint at {input_base_path}.')
					# Load weights
					if model_size == "7B":
										# Not sharded
										# (The sharded implementation would also work, but this is simpler.)
										a__          =  torch.load(os.path.join(lowerCamelCase_ ,     '''consolidated.00.pth''') ,     map_location='''cpu''')
					else:
										# Sharded
										a__          =  [
										    torch.load(os.path.join(lowerCamelCase_ ,     f'consolidated.{i:02d}.pth') ,     map_location='''cpu''')
										    for i in range(lowerCamelCase_)
										]
					a__          =  0
					a__          =  {'''weight_map''': {}}
					for layer_i in range(lowerCamelCase_):
										a__          =  f'pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin'
										if model_size == "7B":
															# Unsharded
															a__          =  {
															    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__          =  {
															    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__          =  permute(
															    torch.cat(
															        [
															            loaded[i][f'layers.{layer_i}.attention.wq.weight'].view(lowerCamelCase_ ,     lowerCamelCase_ ,     lowerCamelCase_)
															            for i in range(lowerCamelCase_)
															        ] ,     dim=0 ,     ).reshape(lowerCamelCase_ ,     lowerCamelCase_))
															a__          =  permute(
															    torch.cat(
															        [
															            loaded[i][f'layers.{layer_i}.attention.wk.weight'].view(
															                lowerCamelCase_ ,     lowerCamelCase_ ,     lowerCamelCase_)
															            for i in range(lowerCamelCase_)
															        ] ,     dim=0 ,     ).reshape(lowerCamelCase_ ,     lowerCamelCase_) ,     lowerCamelCase_ ,     lowerCamelCase_ ,     lowerCamelCase_ ,     )
															a__          =  torch.cat(
															    [
															        loaded[i][f'layers.{layer_i}.attention.wv.weight'].view(
															            lowerCamelCase_ ,     lowerCamelCase_ ,     lowerCamelCase_)
															        for i in range(lowerCamelCase_)
															    ] ,     dim=0 ,     ).reshape(lowerCamelCase_ ,     lowerCamelCase_)

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

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

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

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

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

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

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

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





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_):
					# Initialize the tokenizer based on the `spm` model
					a__          =  LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
					print(f'Saving a {tokenizer_class.__name__} to {tokenizer_path}.')
					a__          =  tokenizer_class(lowerCamelCase_)
					tokenizer.save_pretrained(lowerCamelCase_)





def        SCREAMING_SNAKE_CASE       (      ):
					a__          =  argparse.ArgumentParser()
					parser.add_argument(
					    '''--input_dir''' ,     help='''Location of LLaMA weights, which contains tokenizer.model and model folders''' ,     )
					parser.add_argument(
					    '''--model_size''' ,     choices=['''7B''', '''7Bf''', '''13B''', '''13Bf''', '''30B''', '''65B''', '''70B''', '''70Bf''', '''tokenizer_only'''] ,     )
					parser.add_argument(
					    '''--output_dir''' ,     help='''Location to write HF model and tokenizer''' ,     )
					parser.add_argument('''--safe_serialization''' ,     type=lowerCamelCase_ ,     help='''Whether or not to save using `safetensors`.''')
					a__          =  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__          =  os.path.join(args.input_dir ,     '''tokenizer.model''')
					write_tokenizer(args.output_dir ,     lowerCamelCase_)


if __name__ == "__main__":
	main()


