





'''simple docstring'''



import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import (
		    FocalNetBackbone,
		    FocalNetForImageClassification,
		    FocalNetForMaskedImageModeling,
		    FocalNetModel,
		)
		from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor



class lowerCamelCase			:



			'''simple docstring'''






			def __init__(				self	:						Optional[Any]		,      lowerCAmelCase_	:						Union[str, Any]		,      lowerCAmelCase_	:						Optional[Any]=13		,      lowerCAmelCase_	:						str=32		,      lowerCAmelCase_	:						Any=2		,      lowerCAmelCase_	:						Union[str, Any]=3		,      lowerCAmelCase_	:						int=16		,      lowerCAmelCase_	:						Optional[int]=[32, 64, 1_28]		,      lowerCAmelCase_	:						str=[1, 2, 1]		,      lowerCAmelCase_	:						str=[2, 2, 4]		,      lowerCAmelCase_	:						List[str]=2		,      lowerCAmelCase_	:						List[str]=2.0		,      lowerCAmelCase_	:						List[Any]=True		,      lowerCAmelCase_	:						Tuple=0.0		,      lowerCAmelCase_	:						Optional[Any]=0.0		,      lowerCAmelCase_	:						int=0.1		,      lowerCAmelCase_	:						Optional[int]="gelu"		,      lowerCAmelCase_	:						List[str]=False		,      lowerCAmelCase_	:						Optional[int]=True		,      lowerCAmelCase_	:						Optional[int]=0.02		,      lowerCAmelCase_	:						List[str]=1e-5		,      lowerCAmelCase_	:						List[Any]=True		,      lowerCAmelCase_	:						int=None		,      lowerCAmelCase_	:						List[Any]=True		,      lowerCAmelCase_	:						Optional[Any]=10		,      lowerCAmelCase_	:						Union[str, Any]=8		,      lowerCAmelCase_	:						Optional[Any]=["stage1", "stage2"]		,      lowerCAmelCase_	:						List[Any]=[1, 2]		,      )      ->		int:





									'''simple docstring'''
									A__					:       Tuple     =parent
									A__					:       str     =batch_size
									A__					:       Any     =image_size
									A__					:       List[Any]     =patch_size
									A__					:       List[Any]     =num_channels
									A__					:       List[str]     =embed_dim
									A__					:       str     =hidden_sizes
									A__					:       Any     =depths
									A__					:       List[str]     =num_heads
									A__					:       Any     =window_size
									A__					:       List[str]     =mlp_ratio
									A__					:       Optional[int]     =qkv_bias
									A__					:       Any     =hidden_dropout_prob
									A__					:       List[str]     =attention_probs_dropout_prob
									A__					:       str     =drop_path_rate
									A__					:       Optional[Any]     =hidden_act
									A__					:       Optional[int]     =use_absolute_embeddings
									A__					:       List[str]     =patch_norm
									A__					:       int     =layer_norm_eps
									A__					:       Optional[Any]     =initializer_range
									A__					:       List[str]     =is_training
									A__					:       Dict     =scope
									A__					:       Optional[Any]     =use_labels
									A__					:       Union[str, Any]     =type_sequence_label_size
									A__					:       Optional[int]     =encoder_stride
									A__					:       str     =out_features
									A__					:       Optional[int]     =out_indices


			def lowercase__						(				self	:						Dict			)      ->		Optional[Any]:





									'''simple docstring'''
									A__					:       Tuple     =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

									A__					:       Dict     =None
									if self.use_labels:
															A__					:       List[Any]     =ids_tensor([self.batch_size]		,      self.type_sequence_label_size			)

									A__					:       List[str]     =self.get_config()

									return config, pixel_values, labels


			def lowercase__						(				self	:						Optional[Any]			)      ->		Dict:





									'''simple docstring'''
									return FocalNetConfig(
									    image_size=self.image_size		,      patch_size=self.patch_size		,      num_channels=self.num_channels		,      embed_dim=self.embed_dim		,      hidden_sizes=self.hidden_sizes		,      depths=self.depths		,      num_heads=self.num_heads		,      window_size=self.window_size		,      mlp_ratio=self.mlp_ratio		,      qkv_bias=self.qkv_bias		,      hidden_dropout_prob=self.hidden_dropout_prob		,      attention_probs_dropout_prob=self.attention_probs_dropout_prob		,      drop_path_rate=self.drop_path_rate		,      hidden_act=self.hidden_act		,      use_absolute_embeddings=self.use_absolute_embeddings		,      path_norm=self.patch_norm		,      layer_norm_eps=self.layer_norm_eps		,      initializer_range=self.initializer_range		,      encoder_stride=self.encoder_stride		,      out_features=self.out_features		,      out_indices=self.out_indices		,      )


			def lowercase__						(				self	:						Dict		,      lowerCAmelCase_	:						int		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						str			)      ->		List[Any]:





									'''simple docstring'''
									A__					:       int     =FocalNetModel(config=snake_case_			)
									model.to(snake_case_			)
									model.eval()
									A__					:       Dict     =model(snake_case_			)

									A__					:       Union[str, Any]     =((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths			) - 1))
									A__					:       Dict     =int(config.embed_dim * 2 ** (len(config.depths			) - 1)			)

									self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, expected_seq_len, expected_dim)			)


			def lowercase__						(				self	:						Tuple		,      lowerCAmelCase_	:						Any		,      lowerCAmelCase_	:						Union[str, Any]		,      lowerCAmelCase_	:						Optional[int]			)      ->		List[Any]:





									'''simple docstring'''
									A__					:       List[str]     =FocalNetBackbone(config=snake_case_			)
									model.to(snake_case_			)
									model.eval()
									A__					:       List[str]     =model(snake_case_			)

									# verify feature maps
									self.parent.assertEqual(len(result.feature_maps			)		,      len(config.out_features			)			)
									self.parent.assertListEqual(list(result.feature_maps[0].shape			)		,      [self.batch_size, self.image_size, 8, 8]			)

									# verify channels
									self.parent.assertEqual(len(model.channels			)		,      len(config.out_features			)			)
									self.parent.assertListEqual(model.channels		,      config.hidden_sizes[:-1]			)

									# verify backbone works with out_features=None
									A__					:       Union[str, Any]     =None
									A__					:       Tuple     =FocalNetBackbone(config=snake_case_			)
									model.to(snake_case_			)
									model.eval()
									A__					:       int     =model(snake_case_			)

									# verify feature maps
									self.parent.assertEqual(len(result.feature_maps			)		,      1			)
									self.parent.assertListEqual(list(result.feature_maps[0].shape			)		,      [self.batch_size, self.image_size * 2, 4, 4]			)

									# verify channels
									self.parent.assertEqual(len(model.channels			)		,      1			)
									self.parent.assertListEqual(model.channels		,      [config.hidden_sizes[-1]]			)


			def lowercase__						(				self	:						Optional[Any]		,      lowerCAmelCase_	:						Optional[int]		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						Optional[Any]			)      ->		int:





									'''simple docstring'''
									A__					:       List[str]     =FocalNetForMaskedImageModeling(config=snake_case_			)
									model.to(snake_case_			)
									model.eval()
									A__					:       int     =model(snake_case_			)
									self.parent.assertEqual(
									    result.reconstruction.shape		,      (self.batch_size, self.num_channels, self.image_size, self.image_size)			)

									# test greyscale images
									A__					:       str     =1
									A__					:       Optional[Any]     =FocalNetForMaskedImageModeling(snake_case_			)
									model.to(snake_case_			)
									model.eval()

									A__					:       Any     =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
									A__					:       Union[str, Any]     =model(snake_case_			)
									self.parent.assertEqual(result.reconstruction.shape		,      (self.batch_size, 1, self.image_size, self.image_size)			)


			def lowercase__						(				self	:						Union[str, Any]		,      lowerCAmelCase_	:						List[Any]		,      lowerCAmelCase_	:						Any		,      lowerCAmelCase_	:						int			)      ->		Optional[Any]:





									'''simple docstring'''
									A__					:       Dict     =self.type_sequence_label_size
									A__					:       Optional[Any]     =FocalNetForImageClassification(snake_case_			)
									model.to(snake_case_			)
									model.eval()
									A__					:       Union[str, Any]     =model(snake_case_		,      labels=snake_case_			)
									self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.type_sequence_label_size)			)

									# test greyscale images
									A__					:       Optional[int]     =1
									A__					:       str     =FocalNetForImageClassification(snake_case_			)
									model.to(snake_case_			)
									model.eval()

									A__					:       str     =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
									A__					:       List[Any]     =model(snake_case_			)
									self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.type_sequence_label_size)			)



			def lowercase__						(				self	:						List[Any]			)      ->		int:





									'''simple docstring'''
									A__					:       Any     =self.prepare_config_and_inputs()

									A__					:       Any     =config_and_inputs
									A__					:       Tuple     ={'''pixel_values''': pixel_values}
									return config, inputs_dict



@require_torch
class lowerCamelCase			(  __lowercase	,						__lowercase	,						unittest.TestCase						):



			'''simple docstring'''






			__snake_case        =	(
			    (
			        FocalNetModel,
			        FocalNetForImageClassification,
			        FocalNetForMaskedImageModeling,
			        FocalNetBackbone,
			    )
			    if is_torch_available()
			    else ()
			)
			__snake_case        =	(
			    {"feature-extraction": FocalNetModel, "image-classification": FocalNetForImageClassification}
			    if is_torch_available()
			    else {}
			)
			__snake_case        =	False

			__snake_case        =	False
			__snake_case        =	False
			__snake_case        =	False
			__snake_case        =	False


			def lowercase__						(				self	:						List[Any]			)      ->		Any:





									'''simple docstring'''
									A__					:       Union[str, Any]     =FocalNetModelTester(self			)
									A__					:       Dict     =ConfigTester(self		,      config_class=snake_case_		,      embed_dim=37		,      has_text_modality=snake_case_			)


			def lowercase__						(				self	:						Any			)      ->		List[Any]:





									'''simple docstring'''
									self.create_and_test_config_common_properties()
									self.config_tester.create_and_test_config_to_json_string()
									self.config_tester.create_and_test_config_to_json_file()
									self.config_tester.create_and_test_config_from_and_save_pretrained()
									self.config_tester.create_and_test_config_with_num_labels()
									self.config_tester.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()


			def lowercase__						(				self	:						Dict			)      ->		Optional[Any]:





									'''simple docstring'''
									return


			def lowercase__						(				self	:						Dict			)      ->		Optional[int]:





									'''simple docstring'''
									A__					:       Union[str, Any]     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*snake_case_			)


			def lowercase__						(				self	:						Tuple			)      ->		List[str]:





									'''simple docstring'''
									A__					:       Dict     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_backbone(*snake_case_			)


			def lowercase__						(				self	:						Any			)      ->		List[str]:





									'''simple docstring'''
									A__					:       Dict     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_image_modeling(*snake_case_			)


			def lowercase__						(				self	:						Optional[int]			)      ->		List[Any]:





									'''simple docstring'''
									A__					:       Union[str, Any]     =self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*snake_case_			)


			@unittest.skip(reason="""FocalNet does not use inputs_embeds"""			)
			def lowercase__						(				self	:						Optional[Any]			)      ->		Dict:





									'''simple docstring'''
									pass


			@unittest.skip(reason="""FocalNet does not use feedforward chunking"""			)
			def lowercase__						(				self	:						str			)      ->		str:





									'''simple docstring'''
									pass


			def lowercase__						(				self	:						Tuple			)      ->		str:





									'''simple docstring'''
									A__					:       Dict     =self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes[:-1]:
															A__					:       int     =model_class(snake_case_			)
															self.assertIsInstance(model.get_input_embeddings()		,      (nn.Module)			)
															A__					:       List[Any]     =model.get_output_embeddings()
															self.assertTrue(x is None or isinstance(snake_case_		,      nn.Linear			)			)


			def lowercase__						(				self	:						Optional[int]			)      ->		Tuple:





									'''simple docstring'''
									A__					:       Optional[int]     =self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes[:-1]:
															A__					:       str     =model_class(snake_case_			)
															A__					:       int     =inspect.signature(model.forward			)
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															A__					:       Dict     =[*signature.parameters.keys()]

															A__					:       str     =['''pixel_values''']
															self.assertListEqual(arg_names[:1]		,      snake_case_			)


			def lowercase__						(				self	:						Tuple		,      lowerCAmelCase_	:						Tuple		,      lowerCAmelCase_	:						Dict		,      lowerCAmelCase_	:						Any		,      lowerCAmelCase_	:						Optional[Any]			)      ->		Union[str, Any]:





									'''simple docstring'''
									A__					:       Any     =model_class(snake_case_			)
									model.to(snake_case_			)
									model.eval()

									with torch.no_grad():
															A__					:       Dict     =model(**self._prepare_for_class(snake_case_		,      snake_case_			)			)

									A__					:       Union[str, Any]     =outputs.hidden_states

									A__					:       Dict     =getattr(
									    self.model_tester		,      """expected_num_hidden_layers"""		,      len(self.model_tester.depths			) + 1			)
									self.assertEqual(len(snake_case_			)		,      snake_case_			)

									# FocalNet has a different seq_length
									A__					:       int     =(
									    config.patch_size
									    if isinstance(config.patch_size		,      collections.abc.Iterable			)
									    else (config.patch_size, config.patch_size)
									)

									A__					:       Dict     =(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

									self.assertListEqual(
									    list(hidden_states[0].shape[-2:]			)		,      [num_patches, self.model_tester.embed_dim]		,      )

									A__					:       Optional[Any]     =outputs.reshaped_hidden_states
									self.assertEqual(len(snake_case_			)		,      snake_case_			)

									A__					:       List[Any]     =reshaped_hidden_states[0].shape
									A__					:       List[str]     =(
									    reshaped_hidden_states[0].view(snake_case_		,      snake_case_		,      height * width			).permute(0		,      2		,      1			)
									)
									self.assertListEqual(
									    list(reshaped_hidden_states.shape[-2:]			)		,      [num_patches, self.model_tester.embed_dim]		,      )


			def lowercase__						(				self	:						Optional[int]			)      ->		Any:





									'''simple docstring'''
									A__					:       Dict     =self.model_tester.prepare_config_and_inputs_for_common()

									A__					:       Dict     =(
									    self.model_tester.image_size
									    if isinstance(self.model_tester.image_size		,      collections.abc.Iterable			)
									    else (self.model_tester.image_size, self.model_tester.image_size)
									)

									for model_class in self.all_model_classes[:-1]:
															A__					:       Any     =True
															self.check_hidden_states_output(snake_case_		,      snake_case_		,      snake_case_		,      snake_case_			)

															# check that output_hidden_states also work using config
															del inputs_dict["output_hidden_states"]
															A__					:       Optional[int]     =True

															self.check_hidden_states_output(snake_case_		,      snake_case_		,      snake_case_		,      snake_case_			)


			def lowercase__						(				self	:						List[Any]			)      ->		int:





									'''simple docstring'''
									A__					:       List[Any]     =self.model_tester.prepare_config_and_inputs_for_common()
									A__					:       List[str]     =3

									A__					:       Optional[int]     =(
									    self.model_tester.image_size
									    if isinstance(self.model_tester.image_size		,      collections.abc.Iterable			)
									    else (self.model_tester.image_size, self.model_tester.image_size)
									)
									A__					:       List[str]     =(
									    config.patch_size
									    if isinstance(config.patch_size		,      collections.abc.Iterable			)
									    else (config.patch_size, config.patch_size)
									)

									A__					:       List[Any]     =image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
									A__					:       Union[str, Any]     =image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

									for model_class in self.all_model_classes[:-1]:
															A__					:       int     =True
															self.check_hidden_states_output(snake_case_		,      snake_case_		,      snake_case_		,      (padded_height, padded_width)			)

															# check that output_hidden_states also work using config
															del inputs_dict["output_hidden_states"]
															A__					:       List[str]     =True
															self.check_hidden_states_output(snake_case_		,      snake_case_		,      snake_case_		,      (padded_height, padded_width)			)


			@slow
			def lowercase__						(				self	:						str			)      ->		Optional[Any]:





									'''simple docstring'''
									for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															A__					:       Dict     =FocalNetModel.from_pretrained(snake_case_			)
															self.assertIsNotNone(snake_case_			)



			def lowercase__						(				self	:						List[Any]			)      ->		Optional[Any]:





									'''simple docstring'''
									A__					:       Optional[int]     =self.model_tester.prepare_config_and_inputs_for_common()

									A__					:       Optional[int]     =_config_zero_init(snake_case_			)
									for model_class in self.all_model_classes:
															A__					:       str     =model_class(config=snake_case_			)
															for name, param in model.named_parameters():
																					if "embeddings" not in name and param.requires_grad:
																											self.assertIn(
																											    ((param.data.mean() * 1e9).round() / 1e9).item()		,      [0.0, 1.0]		,      msg=f"Parameter {name} of model {model_class} seems not properly initialized"		,      )



@require_vision
@require_torch
class lowerCamelCase			(  unittest.TestCase						):



			'''simple docstring'''






			@cached_property
			def lowercase__						(				self	:						str			)      ->		Any:





									'''simple docstring'''
									# TODO update organization
									return AutoImageProcessor.from_pretrained("""microsoft/focalnet-tiny"""			) if is_vision_available() else None



			@slow
			def lowercase__						(				self	:						str			)      ->		Optional[Any]:





									'''simple docstring'''
									A__					:       int     =FocalNetForImageClassification.from_pretrained("""microsoft/focalnet-tiny"""			).to(snake_case_			)
									A__					:       Optional[Any]     =self.default_image_processor

									A__					:       Dict     =Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
									A__					:       Optional[Any]     =image_processor(images=snake_case_		,      return_tensors="""pt"""			).to(snake_case_			)

									# forward pass
									with torch.no_grad():
															A__					:       Any     =model(**snake_case_			)

									# verify the logits
									A__					:       List[Any]     =torch.Size((1, 10_00)			)
									self.assertEqual(outputs.logits.shape		,      snake_case_			)
									A__					:       Union[str, Any]     =torch.tensor([0.2166, -0.4368, 0.2191]			).to(snake_case_			)
									self.assertTrue(torch.allclose(outputs.logits[0, :3]		,      snake_case_		,      atol=1e-4			)			)
									self.assertTrue(outputs.logits.argmax(dim=-1			).item()		,      2_81			)



@require_torch
class lowerCamelCase			(  __lowercase	,						unittest.TestCase						):



			'''simple docstring'''






			__snake_case        =	(FocalNetBackbone,) if is_torch_available() else ()
			__snake_case        =	FocalNetConfig

			__snake_case        =	False


			def lowercase__						(				self	:						Tuple			)      ->		Any:





									'''simple docstring'''
									A__					:       Union[str, Any]     =FocalNetModelTester(self			)




'''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'
		)
		__snake_case			:     Optional[Any]										=      None
__snake_case			:     Optional[Any]										=      {
    '7B': 1_1008,
    '13B': 1_3824,
    '30B': 1_7920,
    '65B': 2_2016,
    '70B': 2_8672,
}
__snake_case			:     Union[str, Any]										=      {
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}
def 						__lowerCamelCase	(       __snake_case					:		Optional[Any],		__snake_case					:		str=1,		__snake_case					:		Tuple=256	)	->							str:






						"""simple docstring"""





						return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3	)	) + multiple_of - 1) // multiple_of)
def 						__lowerCamelCase	(       __snake_case					:		Tuple	)	->							Tuple:






						"""simple docstring"""





						with open(__snake_case,		"""r"""	) as f:
												return json.load(__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		Optional[int],		__snake_case					:		Tuple	)	->							Dict:






						"""simple docstring"""





						with open(__snake_case,		"""w"""	) as f:
												json.dump(__snake_case,		__snake_case	)
def 						__lowerCamelCase	(       __snake_case					:		List[Any],		__snake_case					:		Any,		__snake_case					:		Any,		__snake_case					:		Tuple=True	)	->							Any:






						"""simple docstring"""





						os.makedirs(__snake_case,		exist_ok=__snake_case	)
						A__					:       List[Any]     =os.path.join(__snake_case,		"""tmp"""	)
						os.makedirs(__snake_case,		exist_ok=__snake_case	)

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

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

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

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

												A__					:       List[Any]     =inv_freq
												for k, v in state_dict.items():
																		A__					:       Optional[Any]     =filename
																		param_count += v.numel()
												torch.save(__snake_case,		os.path.join(__snake_case,		__snake_case	)	)

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

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






						"""simple docstring"""





						A__					:       List[Any]     =LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
						print(f"Saving a {tokenizer_class.__name__} to {tokenizer_path}."	)
						A__					:       List[str]     =tokenizer_class(__snake_case	)
						tokenizer.save_pretrained(__snake_case	)
def 						__lowerCamelCase	(       )	->							Any:






						"""simple docstring"""





						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=__snake_case,		help="""Whether or not to save using `safetensors`."""	)
						A__					:       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,		__snake_case	)


if __name__ == "__main__":
		main()

