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

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


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







def  SCREAMING_SNAKE_CASE_		(    __A				:    str						,					__A				:    Any=1						,					__A				:    Optional[Any]=2_56       )			->	Union[str, Any]:

       """simple docstring"""

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







def  SCREAMING_SNAKE_CASE_		(    __A				:    int       )			->	List[str]:

       """simple docstring"""

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







def  SCREAMING_SNAKE_CASE_		(    __A				:    Dict						,					__A				:    Optional[Any]       )			->	Optional[Any]:

       """simple docstring"""

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







def  SCREAMING_SNAKE_CASE_		(    __A				:    Tuple						,					__A				:    Any						,					__A				:    Any						,					__A				:    str=True       )			->	Optional[Any]:

       """simple docstring"""

       os.makedirs(__snake_case						,					exist_ok=__snake_case       )
       a_      :	Union[str, Any]							=						os.path.join(__snake_case						,					'tmp'       )
       os.makedirs(__snake_case						,					exist_ok=__snake_case       )

       a_      :	str							=						read_json(os.path.join(__snake_case						,					'params.json'       )       )
       a_      :	int							=						NUM_SHARDS[model_size]
       a_      :	List[Any]							=						params['n_layers']
       a_      :	Optional[Any]							=						params['n_heads']
       a_      :	Any							=						n_heads // num_shards
       a_      :	Union[str, Any]							=						params['dim']
       a_      :	int							=						dim // n_heads
       a_      :	int							=						10000.0
       a_      :	int							=						1.0 / (base ** (torch.arange(0						,					__snake_case						,					2       ).float() / dims_per_head))

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

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

                     a_      :	Tuple							=						{
                         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_      :	Dict							=						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_      :	Optional[Any]							=						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_      :	Dict							=						torch.cat(
                         [loaded[i][F"""layers.{layer_i}.attention.wo.weight"""] for i in range(__snake_case       )]						,					dim=1       )
                     a_      :	Any							=						torch.cat(
                         [loaded[i][F"""layers.{layer_i}.feed_forward.w1.weight"""] for i in range(__snake_case       )]						,					dim=0       )
                     a_      :	List[Any]							=						torch.cat(
                         [loaded[i][F"""layers.{layer_i}.feed_forward.w2.weight"""] for i in range(__snake_case       )]						,					dim=1       )
                     a_      :	Dict							=						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_      :	str							=						filename
                     param_count += v.numel()
              torch.save(__snake_case						,					os.path.join(__snake_case						,					__snake_case       )       )

       a_      :	Any							=						F"""pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"""
       if model_size == "7B":
              # Unsharded
              a_      :	List[str]							=						{
                  'model.embed_tokens.weight': loaded['tok_embeddings.weight'],
                  'model.norm.weight': loaded['norm.weight'],
                  'lm_head.weight': loaded['output.weight'],
              }
       else:
              a_      :	Tuple							=						{
                  '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_      :	int							=						{'total_size': param_count * 2}
       write_json(__snake_case						,					os.path.join(__snake_case						,					'pytorch_model.bin.index.json'       )       )
       a_      :	int							=						params['ffn_dim_multiplier'] if 'ffn_dim_multiplier' in params else 1
       a_      :	Optional[Any]							=						params['multiple_of'] if 'multiple_of' in params else 2_56
       a_      :	Any							=						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_      :	Dict							=						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  SCREAMING_SNAKE_CASE_		(    __A				:    Optional[Any]						,					__A				:    Optional[Any]       )			->	Any:

       """simple docstring"""

       a_      :	List[Any]							=						LlamaTokenizer if LlamaTokenizerFast is None else LlamaTokenizerFast
       print(F"""Saving a {tokenizer_class.__name__} to {tokenizer_path}."""       )
       a_      :	Optional[Any]							=						tokenizer_class(__snake_case       )
       tokenizer.save_pretrained(__snake_case       )







def  SCREAMING_SNAKE_CASE_		(    )			->	int:

       """simple docstring"""

       a_      :	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=__snake_case						,					help='Whether or not to save using `safetensors`.'       )
       a_      :	Tuple							=						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_      :	Any							=						os.path.join(args.input_dir						,					'tokenizer.model'       )
       write_tokenizer(args.output_dir						,					__snake_case       )


if __name__ == "__main__":
   main()






from typing import List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_   :				Dict						      = logging.get_logger(__name__)

UpperCAmelCase_   :				int						      = {
    'huggingface/autoformer-tourism-monthly': 'https://huggingface.co/huggingface/autoformer-tourism-monthly/resolve/main/config.json',
}
class 				SCREAMING_SNAKE_CASE__						(       lowercase__    ):
      snake_case__       :  Union[str, Any]       =		'''autoformer'''
      snake_case__       :  Any       =		{
          '''hidden_size''': '''d_model''',
          '''num_attention_heads''': '''encoder_attention_heads''',
          '''num_hidden_layers''': '''encoder_layers''',
      }
      def __init__( self       : Optional[int]  ,							SCREAMING_SNAKE_CASE__       : Optional[int] = None  ,							SCREAMING_SNAKE_CASE__       : Optional[int] = None  ,							SCREAMING_SNAKE_CASE__       : str = "student_t"  ,							SCREAMING_SNAKE_CASE__       : str = "nll"  ,							SCREAMING_SNAKE_CASE__       : int = 1  ,							SCREAMING_SNAKE_CASE__       : List[int] = [1, 2, 3, 4, 5, 6, 7]  ,							SCREAMING_SNAKE_CASE__       : bool = True  ,							SCREAMING_SNAKE_CASE__       : int = 0  ,							SCREAMING_SNAKE_CASE__       : int = 0  ,							SCREAMING_SNAKE_CASE__       : int = 0  ,							SCREAMING_SNAKE_CASE__       : int = 0  ,							SCREAMING_SNAKE_CASE__       : Optional[List[int]] = None  ,							SCREAMING_SNAKE_CASE__       : Optional[List[int]] = None  ,							SCREAMING_SNAKE_CASE__       : int = 6_4  ,							SCREAMING_SNAKE_CASE__       : int = 2  ,							SCREAMING_SNAKE_CASE__       : int = 2  ,							SCREAMING_SNAKE_CASE__       : int = 2  ,							SCREAMING_SNAKE_CASE__       : int = 2  ,							SCREAMING_SNAKE_CASE__       : int = 3_2  ,							SCREAMING_SNAKE_CASE__       : int = 3_2  ,							SCREAMING_SNAKE_CASE__       : str = "gelu"  ,							SCREAMING_SNAKE_CASE__       : float = 0.1  ,							SCREAMING_SNAKE_CASE__       : float = 0.1  ,							SCREAMING_SNAKE_CASE__       : float = 0.1  ,							SCREAMING_SNAKE_CASE__       : float = 0.1  ,							SCREAMING_SNAKE_CASE__       : float = 0.1  ,							SCREAMING_SNAKE_CASE__       : int = 1_0_0  ,							SCREAMING_SNAKE_CASE__       : float = 0.02  ,							SCREAMING_SNAKE_CASE__       : bool = True  ,							SCREAMING_SNAKE_CASE__       : List[str]=True  ,							SCREAMING_SNAKE_CASE__       : int = 1_0  ,							SCREAMING_SNAKE_CASE__       : int = 2_5  ,							SCREAMING_SNAKE_CASE__       : int = 3  ,							**SCREAMING_SNAKE_CASE__       : int  ,							)	->    Union[str, Any]:
             # time series specific configuration
             a_      :	Optional[int]							=						prediction_length
             a_      :	Tuple							=						context_length if context_length is not None else prediction_length
             a_      :	Tuple							=						distribution_output
             a_      :	int							=						loss
             a_      :	Tuple							=						input_size
             a_      :	int							=						num_time_features
             a_      :	Dict							=						lags_sequence
             a_      :	List[str]							=						scaling
             a_      :	Tuple							=						num_dynamic_real_features
             a_      :	Dict							=						num_static_real_features
             a_      :	Dict							=						num_static_categorical_features
             if cardinality is not None and num_static_categorical_features > 0:
                    if len(SCREAMING_SNAKE_CASE__  ) != num_static_categorical_features:
                           raise ValueError(
                               'The cardinality should be a list of the same length as `num_static_categorical_features`'  )
                    a_      :	Any							=						cardinality
             else:
                    a_      :	Tuple							=						[0]
             if embedding_dimension is not None and num_static_categorical_features > 0:
                    if len(SCREAMING_SNAKE_CASE__  ) != num_static_categorical_features:
                           raise ValueError(
                               'The embedding dimension should be a list of the same length as `num_static_categorical_features`'  )
                    a_      :	Tuple							=						embedding_dimension
             else:
                    a_      :	Optional[Any]							=						[min(5_0  ,							(cat + 1) // 2  ) for cat in self.cardinality]
             a_      :	Optional[Any]							=						num_parallel_samples

             # Transformer architecture configuration
             a_      :	Dict							=						input_size * len(self.lags_sequence  ) + self._number_of_features
             a_      :	Tuple							=						d_model
             a_      :	int							=						encoder_attention_heads
             a_      :	Optional[int]							=						decoder_attention_heads
             a_      :	str							=						encoder_ffn_dim
             a_      :	List[str]							=						decoder_ffn_dim
             a_      :	Any							=						encoder_layers
             a_      :	Any							=						decoder_layers

             a_      :	List[Any]							=						dropout
             a_      :	Tuple							=						attention_dropout
             a_      :	Union[str, Any]							=						activation_dropout
             a_      :	List[Any]							=						encoder_layerdrop
             a_      :	List[str]							=						decoder_layerdrop

             a_      :	Optional[Any]							=						activation_function
             a_      :	Union[str, Any]							=						init_std

             a_      :	List[str]							=						use_cache

             # Autoformer
             a_      :	Tuple							=						label_length
             a_      :	List[Any]							=						moving_average
             a_      :	Optional[Any]							=						autocorrelation_factor

             super().__init__(is_encoder_decoder=SCREAMING_SNAKE_CASE__  ,							**SCREAMING_SNAKE_CASE__  )

      @property
      def  SCREAMING_SNAKE_CASE							( self       : List[Any]  )	->    int:
             return (
                 sum(self.embedding_dimension  )
                 + self.num_dynamic_real_features
                 + self.num_time_features
                 + self.num_static_real_features
                 + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
             )


