





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


lowerCAmelCase_    =    logging.get_logger(__name__)

# General docstring
lowerCAmelCase_    =    '''ResNetConfig'''

# Base docstring
lowerCAmelCase_    =    '''microsoft/resnet-50'''
lowerCAmelCase_    =    [1, 2_0_4_8, 7, 7]

# Image classification docstring
lowerCAmelCase_    =    '''microsoft/resnet-50'''
lowerCAmelCase_    =    '''tiger cat'''

lowerCAmelCase_    =    [
    '''microsoft/resnet-50''',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]



class  __lowerCAmelCase	(      nn.Module					):






			def __init__(self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ = 3 ,							__magic_name__ = 1 ,							__magic_name__ = "relu"   )					->       str:


							'''simple docstring'''



							super().__init__()
							snake_case_						:   int           =  nn.Convad(
							    _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							kernel_size=_SCREAMING_SNAKE_CASE ,							stride=_SCREAMING_SNAKE_CASE ,							padding=kernel_size // 2 ,							bias=_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Union[str, Any]           =  nn.BatchNormad(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Dict           =  ACTaFN[activation] if activation is not None else nn.Identity()






			def       lowerCamelCase  (self ,							__magic_name__   )					->       Tensor:


							'''simple docstring'''



							snake_case_						:   List[str]           =  self.convolution(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Optional[Any]           =  self.normalization(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Dict           =  self.activation(_SCREAMING_SNAKE_CASE   )
							return hidden_state




class  __lowerCAmelCase	(      nn.Module					):






			def __init__(self ,							__magic_name__   )					->       Optional[Any]:


							'''simple docstring'''



							super().__init__()
							snake_case_						:   Dict           =  ResNetConvLayer(
							    config.num_channels ,							config.embedding_size ,							kernel_size=7 ,							stride=2 ,							activation=config.hidden_act   )
							snake_case_						:   str           =  nn.MaxPoolad(kernel_size=3 ,							stride=2 ,							padding=1   )
							snake_case_						:   List[str]           =  config.num_channels






			def       lowerCamelCase  (self ,							__magic_name__   )					->       Tensor:


							'''simple docstring'''



							snake_case_						:   Optional[Any]           =  pixel_values.shape[1]
							if num_channels != self.num_channels:
											raise ValueError(
											    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.'''   )
							snake_case_						:   Dict           =  self.embedder(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Dict           =  self.pooler(_SCREAMING_SNAKE_CASE   )
							return embedding




class  __lowerCAmelCase	(      nn.Module					):






			def __init__(self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ = 2   )					->       List[Any]:


							'''simple docstring'''



							super().__init__()
							snake_case_						:   Union[str, Any]           =  nn.Convad(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							kernel_size=1 ,							stride=_SCREAMING_SNAKE_CASE ,							bias=_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Union[str, Any]           =  nn.BatchNormad(_SCREAMING_SNAKE_CASE   )






			def       lowerCamelCase  (self ,							__magic_name__   )					->       Tensor:


							'''simple docstring'''



							snake_case_						:   List[str]           =  self.convolution(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Dict           =  self.normalization(_SCREAMING_SNAKE_CASE   )
							return hidden_state




class  __lowerCAmelCase	(      nn.Module					):






			def __init__(self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ = 1 ,							__magic_name__ = "relu"   )					->       Tuple:


							'''simple docstring'''



							super().__init__()
							snake_case_						:   Optional[Any]           =  in_channels != out_channels or stride != 1
							snake_case_						:   Tuple           =  (
							    ResNetShortCut(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							stride=_SCREAMING_SNAKE_CASE   ) if should_apply_shortcut else nn.Identity()
							)
							snake_case_						:   Optional[int]           =  nn.Sequential(
							    ResNetConvLayer(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							stride=_SCREAMING_SNAKE_CASE   ) ,							ResNetConvLayer(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							activation=_SCREAMING_SNAKE_CASE   ) ,							)
							snake_case_						:   Any           =  ACTaFN[activation]






			def       lowerCamelCase  (self ,							__magic_name__   )					->       int:


							'''simple docstring'''



							snake_case_						:   Union[str, Any]           =  hidden_state
							snake_case_						:   Union[str, Any]           =  self.layer(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Dict           =  self.shortcut(_SCREAMING_SNAKE_CASE   )
							hidden_state += residual
							snake_case_						:   Dict           =  self.activation(_SCREAMING_SNAKE_CASE   )
							return hidden_state




class  __lowerCAmelCase	(      nn.Module					):






			def __init__(self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ = 1 ,							__magic_name__ = "relu" ,							__magic_name__ = 4   )					->       Dict:


							'''simple docstring'''



							super().__init__()
							snake_case_						:   Optional[Any]           =  in_channels != out_channels or stride != 1
							snake_case_						:   Tuple           =  out_channels // reduction
							snake_case_						:   Tuple           =  (
							    ResNetShortCut(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							stride=_SCREAMING_SNAKE_CASE   ) if should_apply_shortcut else nn.Identity()
							)
							snake_case_						:   List[str]           =  nn.Sequential(
							    ResNetConvLayer(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							kernel_size=1   ) ,							ResNetConvLayer(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							stride=_SCREAMING_SNAKE_CASE   ) ,							ResNetConvLayer(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							kernel_size=1 ,							activation=_SCREAMING_SNAKE_CASE   ) ,							)
							snake_case_						:   str           =  ACTaFN[activation]






			def       lowerCamelCase  (self ,							__magic_name__   )					->       int:


							'''simple docstring'''



							snake_case_						:   int           =  hidden_state
							snake_case_						:   Optional[Any]           =  self.layer(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   List[str]           =  self.shortcut(_SCREAMING_SNAKE_CASE   )
							hidden_state += residual
							snake_case_						:   Dict           =  self.activation(_SCREAMING_SNAKE_CASE   )
							return hidden_state




class  __lowerCAmelCase	(      nn.Module					):






			def __init__(self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ ,							__magic_name__ = 2 ,							__magic_name__ = 2 ,							)					->       Tuple:


							'''simple docstring'''



							super().__init__()

							snake_case_						:   List[str]           =  ResNetBottleNeckLayer if config.layer_type == "bottleneck" else ResNetBasicLayer

							snake_case_						:   List[Any]           =  nn.Sequential(
							    # downsampling is done in the first layer with stride of 2
							    layer(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							stride=_SCREAMING_SNAKE_CASE ,							activation=config.hidden_act   ) ,							*[layer(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							activation=config.hidden_act   ) for _ in range(depth - 1   )] ,							)






			def       lowerCamelCase  (self ,							__magic_name__   )					->       Tensor:


							'''simple docstring'''



							snake_case_						:   str           =  input
							for layer in self.layers:
											snake_case_						:   Optional[Any]           =  layer(_SCREAMING_SNAKE_CASE   )
							return hidden_state




class  __lowerCAmelCase	(      nn.Module					):






			def __init__(self ,							__magic_name__   )					->       Dict:


							'''simple docstring'''



							super().__init__()
							snake_case_						:   Optional[Any]           =  nn.ModuleList([]   )
							# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
							self.stages.append(
							    ResNetStage(
							        _SCREAMING_SNAKE_CASE ,							config.embedding_size ,							config.hidden_sizes[0] ,							stride=2 if config.downsample_in_first_stage else 1 ,							depth=config.depths[0] ,							)   )
							snake_case_						:   str           =  zip(config.hidden_sizes ,							config.hidden_sizes[1:]   )
							for (in_channels, out_channels), depth in zip(_SCREAMING_SNAKE_CASE ,							config.depths[1:]   ):
											self.stages.append(ResNetStage(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							depth=_SCREAMING_SNAKE_CASE   )   )






			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__ = False ,							__magic_name__ = True   )					->       BaseModelOutputWithNoAttention:


							'''simple docstring'''



							snake_case_						:   int           =  () if output_hidden_states else None

							for stage_module in self.stages:
											if output_hidden_states:
															snake_case_						:   Tuple           =  hidden_states + (hidden_state,)

											snake_case_						:   Optional[Any]           =  stage_module(_SCREAMING_SNAKE_CASE   )

							if output_hidden_states:
											snake_case_						:   Optional[int]           =  hidden_states + (hidden_state,)

							if not return_dict:
											return tuple(v for v in [hidden_state, hidden_states] if v is not None   )

							return BaseModelOutputWithNoAttention(
							    last_hidden_state=_SCREAMING_SNAKE_CASE ,							hidden_states=_SCREAMING_SNAKE_CASE ,							)




class  __lowerCAmelCase	(      SCREAMING_SNAKE_CASE__					):
			lowerCamelCase_					:					Tuple	      =				ResNetConfig
			lowerCamelCase_					:					List[Any]	      =				'resnet'
			lowerCamelCase_					:					Dict	      =				'pixel_values'
			lowerCamelCase_					:					Optional[int]	      =				True






			def       lowerCamelCase  (self ,							__magic_name__   )					->       Optional[int]:


							'''simple docstring'''



							if isinstance(_SCREAMING_SNAKE_CASE ,							nn.Convad   ):
											nn.init.kaiming_normal_(module.weight ,							mode='''fan_out''' ,							nonlinearity='''relu'''   )
							elif isinstance(_SCREAMING_SNAKE_CASE ,							(nn.BatchNormad, nn.GroupNorm)   ):
											nn.init.constant_(module.weight ,							1   )
											nn.init.constant_(module.bias ,							0   )






			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__=False   )					->       Optional[int]:


							'''simple docstring'''



							if isinstance(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE   ):
											snake_case_						:   List[Any]           =  value




lowerCAmelCase_    =    r'''
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

lowerCAmelCase_    =    r'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
'''



@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.''',					SCREAMING_SNAKE_CASE__,					)
class  __lowerCAmelCase	(      SCREAMING_SNAKE_CASE__					):






			def __init__(self ,							__magic_name__   )					->       Optional[Any]:


							'''simple docstring'''



							super().__init__(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Tuple           =  config
							snake_case_						:   Optional[int]           =  ResNetEmbeddings(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Optional[int]           =  ResNetEncoder(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Optional[Any]           =  nn.AdaptiveAvgPoolad((1, 1)   )
							# Initialize weights and apply final processing
							self.post_init()






			@add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE   )
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC ,							output_type=_SCREAMING_SNAKE_CASE ,							config_class=_CONFIG_FOR_DOC ,							modality='''vision''' ,							expected_output=_EXPECTED_OUTPUT_SHAPE ,							)
			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__ = None ,							__magic_name__ = None   )					->       BaseModelOutputWithPoolingAndNoAttention:


							'''simple docstring'''



							snake_case_						:   Optional[Any]           =  (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							snake_case_						:   Tuple           =  return_dict if return_dict is not None else self.config.use_return_dict

							snake_case_						:   Any           =  self.embedder(_SCREAMING_SNAKE_CASE   )

							snake_case_						:   Tuple           =  self.encoder(
							    _SCREAMING_SNAKE_CASE ,							output_hidden_states=_SCREAMING_SNAKE_CASE ,							return_dict=_SCREAMING_SNAKE_CASE   )

							snake_case_						:   Any           =  encoder_outputs[0]

							snake_case_						:   str           =  self.pooler(_SCREAMING_SNAKE_CASE   )

							if not return_dict:
											return (last_hidden_state, pooled_output) + encoder_outputs[1:]

							return BaseModelOutputWithPoolingAndNoAttention(
							    last_hidden_state=_SCREAMING_SNAKE_CASE ,							pooler_output=_SCREAMING_SNAKE_CASE ,							hidden_states=encoder_outputs.hidden_states ,							)




@add_start_docstrings(
    '''\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    ''',					SCREAMING_SNAKE_CASE__,					)
class  __lowerCAmelCase	(      SCREAMING_SNAKE_CASE__					):






			def __init__(self ,							__magic_name__   )					->       Dict:


							'''simple docstring'''



							super().__init__(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Union[str, Any]           =  config.num_labels
							snake_case_						:   Union[str, Any]           =  ResNetModel(_SCREAMING_SNAKE_CASE   )
							# classification head
							snake_case_						:   Any           =  nn.Sequential(
							    nn.Flatten() ,							nn.Linear(config.hidden_sizes[-1] ,							config.num_labels   ) if config.num_labels > 0 else nn.Identity() ,							)
							# initialize weights and apply final processing
							self.post_init()






			@add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE   )
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT ,							output_type=_SCREAMING_SNAKE_CASE ,							config_class=_CONFIG_FOR_DOC ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT ,							)
			def       lowerCamelCase  (self ,							__magic_name__ = None ,							__magic_name__ = None ,							__magic_name__ = None ,							__magic_name__ = None ,							)					->       ImageClassifierOutputWithNoAttention:


							'''simple docstring'''



							snake_case_						:   Tuple           =  return_dict if return_dict is not None else self.config.use_return_dict

							snake_case_						:   Any           =  self.resnet(_SCREAMING_SNAKE_CASE ,							output_hidden_states=_SCREAMING_SNAKE_CASE ,							return_dict=_SCREAMING_SNAKE_CASE   )

							snake_case_						:   Dict           =  outputs.pooler_output if return_dict else outputs[1]

							snake_case_						:   str           =  self.classifier(_SCREAMING_SNAKE_CASE   )

							snake_case_						:   Dict           =  None

							if labels is not None:
											if self.config.problem_type is None:
															if self.num_labels == 1:
																			snake_case_						:   Tuple           =  "regression"
															elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																			snake_case_						:   Tuple           =  "single_label_classification"
															else:
																			snake_case_						:   int           =  "multi_label_classification"
											if self.config.problem_type == "regression":
															snake_case_						:   str           =  MSELoss()
															if self.num_labels == 1:
																			snake_case_						:   Union[str, Any]           =  loss_fct(logits.squeeze() ,							labels.squeeze()   )
															else:
																			snake_case_						:   int           =  loss_fct(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE   )
											elif self.config.problem_type == "single_label_classification":
															snake_case_						:   Optional[Any]           =  CrossEntropyLoss()
															snake_case_						:   int           =  loss_fct(logits.view(-1 ,							self.num_labels   ) ,							labels.view(-1   )   )
											elif self.config.problem_type == "multi_label_classification":
															snake_case_						:   Union[str, Any]           =  BCEWithLogitsLoss()
															snake_case_						:   Tuple           =  loss_fct(_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE   )

							if not return_dict:
											snake_case_						:   List[str]           =  (logits,) + outputs[2:]
											return (loss,) + output if loss is not None else output

							return ImageClassifierOutputWithNoAttention(loss=_SCREAMING_SNAKE_CASE ,							logits=_SCREAMING_SNAKE_CASE ,							hidden_states=outputs.hidden_states   )




@add_start_docstrings(
    '''\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    ''',					SCREAMING_SNAKE_CASE__,					)
class  __lowerCAmelCase	(      SCREAMING_SNAKE_CASE__,					SCREAMING_SNAKE_CASE__					):






			def __init__(self ,							__magic_name__   )					->       int:


							'''simple docstring'''



							super().__init__(_SCREAMING_SNAKE_CASE   )
							super()._init_backbone(_SCREAMING_SNAKE_CASE   )

							snake_case_						:   str           =  [config.embedding_size] + config.hidden_sizes
							snake_case_						:   List[Any]           =  ResNetEmbeddings(_SCREAMING_SNAKE_CASE   )
							snake_case_						:   Dict           =  ResNetEncoder(_SCREAMING_SNAKE_CASE   )

							# initialize weights and apply final processing
							self.post_init()






			@add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE   )
			@replace_return_docstrings(output_type=_SCREAMING_SNAKE_CASE ,							config_class=_CONFIG_FOR_DOC   )
			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__ = None ,							__magic_name__ = None   )					->       BackboneOutput:


							'''simple docstring'''



							snake_case_						:   List[str]           =  return_dict if return_dict is not None else self.config.use_return_dict
							snake_case_						:   Dict           =  (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)

							snake_case_						:   Union[str, Any]           =  self.embedder(_SCREAMING_SNAKE_CASE   )

							snake_case_						:   List[Any]           =  self.encoder(_SCREAMING_SNAKE_CASE ,							output_hidden_states=_SCREAMING_SNAKE_CASE ,							return_dict=_SCREAMING_SNAKE_CASE   )

							snake_case_						:   List[Any]           =  outputs.hidden_states

							snake_case_						:   Any           =  ()
							for idx, stage in enumerate(self.stage_names   ):
											if stage in self.out_features:
															feature_maps += (hidden_states[idx],)

							if not return_dict:
											snake_case_						:   Union[str, Any]           =  (feature_maps,)
											if output_hidden_states:
															output += (outputs.hidden_states,)
											return output

							return BackboneOutput(
							    feature_maps=_SCREAMING_SNAKE_CASE ,							hidden_states=outputs.hidden_states if output_hidden_states else None ,							attentions=_SCREAMING_SNAKE_CASE ,							)







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()
