

"""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'
 )
 SCREAMING_SNAKE_CASE_					:	List[str]							  =						None
SCREAMING_SNAKE_CASE_					:	Optional[int]							  =						{
    '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,
}
SCREAMING_SNAKE_CASE_					:	List[Any]							  =						{
    '7B': 1,
    '7Bf': 1,
    '13B': 2,
    '13Bf': 2,
    '30B': 4,
    '65B': 8,
    '70B': 8,
    '70Bf': 8,
}





def  _snake_case     (    UpperCAmelCase_  :					Optional[int]					,   UpperCAmelCase_  :					Optional[int]=1					,   UpperCAmelCase_  :					Any=256			):
     return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3			)			) + multiple_of - 1) // multiple_of)





def  _snake_case     (    UpperCAmelCase_  :					Optional[Any]			):
     with open(UpperCAmelCase_					,   """r"""			) as f:
          return json.load(UpperCAmelCase_			)





def  _snake_case     (    UpperCAmelCase_  :					List[str]					,   UpperCAmelCase_  :					Dict			):
     with open(UpperCAmelCase_					,   """w"""			) as f:
          json.dump(UpperCAmelCase_					,   UpperCAmelCase_			)





def  _snake_case     (    UpperCAmelCase_  :					List[Any]					,   UpperCAmelCase_  :					Optional[Any]					,   UpperCAmelCase_  :					Any					,   UpperCAmelCase_  :					List[Any]=True			):
     os.makedirs(UpperCAmelCase_					,   exist_ok=UpperCAmelCase_			)
     A__										=  os.path.join(UpperCAmelCase_					,   """tmp"""			)
     os.makedirs(UpperCAmelCase_					,   exist_ok=UpperCAmelCase_			)

     A__										=  read_json(os.path.join(UpperCAmelCase_					,   """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__										=  1_00_00.0
     A__										=  1.0 / (base ** (torch.arange(0					,   UpperCAmelCase_					,   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(UpperCAmelCase_  :					Dict					,   UpperCAmelCase_  :					List[str]=n_heads					,   UpperCAmelCase_  :					List[str]=dim					,   UpperCAmelCase_  :					Optional[int]=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.)
          A__										=  torch.load(os.path.join(UpperCAmelCase_					,   """consolidated.00.pth"""			)					,   map_location="""cpu"""			)
     else:
          # Sharded
          A__										=  [
              torch.load(os.path.join(UpperCAmelCase_					,   F"""consolidated.{i:02d}.pth"""			)					,   map_location="""cpu"""			)
              for i in range(UpperCAmelCase_			)
          ]
     A__										=  0
     A__										=  {"""weight_map""": {}}
     for layer_i in range(UpperCAmelCase_			):
          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(UpperCAmelCase_					,   UpperCAmelCase_					,   UpperCAmelCase_			)
                           for i in range(UpperCAmelCase_			)
                       ]					,   dim=0					,   ).reshape(UpperCAmelCase_					,   UpperCAmelCase_			)			)
               A__										=  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_					,   )
               A__										=  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_			)

               A__										=  torch.cat(
                   [loaded[i][F"""layers.{layer_i}.attention.wo.weight"""] for i in range(UpperCAmelCase_			)]					,   dim=1			)
               A__										=  torch.cat(
                   [loaded[i][F"""layers.{layer_i}.feed_forward.w1.weight"""] for i in range(UpperCAmelCase_			)]					,   dim=0			)
               A__										=  torch.cat(
                   [loaded[i][F"""layers.{layer_i}.feed_forward.w2.weight"""] for i in range(UpperCAmelCase_			)]					,   dim=1			)
               A__										=  torch.cat(
                   [loaded[i][F"""layers.{layer_i}.feed_forward.w3.weight"""] for i in range(UpperCAmelCase_			)]					,   dim=0			)

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

     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(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():
          A__										=  filename
          param_count += v.numel()
     torch.save(UpperCAmelCase_					,   os.path.join(UpperCAmelCase_					,   UpperCAmelCase_			)			)

     # Write configs
     A__										=  {"""total_size""": param_count * 2}
     write_json(UpperCAmelCase_					,   os.path.join(UpperCAmelCase_					,   """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=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."""			)
     A__										=  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  _snake_case     (    UpperCAmelCase_  :					List[Any]					,   UpperCAmelCase_  :					List[str]			):
     # 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(UpperCAmelCase_			)
     tokenizer.save_pretrained(UpperCAmelCase_			)





def  _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=UpperCAmelCase_					,   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					,   UpperCAmelCase_			)


if __name__ == "__main__":
 main()









"""simple docstring"""



from __future__ import annotations



class 					a						:



    """simple docstring"""



    def __init__(     self: Any    ,						UpperCamelCase: str    ,						UpperCamelCase: str       ):






         """simple docstring"""


         A__			,				A__										=  text, pattern
         A__			,				A__										=  len(UpperCamelCase       ), len(UpperCamelCase       )

    def 				UpperCamelCase	(     self: Dict    ,						UpperCamelCase: str       ):






         """simple docstring"""


         for i in range(self.patLen - 1    ,						-1    ,						-1       ):
              if char == self.pattern[i]:
                   return i
         return -1

    def 				UpperCamelCase	(     self: str    ,						UpperCamelCase: int       ):






         """simple docstring"""


         for i in range(self.patLen - 1    ,						-1    ,						-1       ):
              if self.pattern[i] != self.text[current_pos + i]:
                   return current_pos + i
         return -1





    def 				UpperCamelCase	(     self: List[Any]       ):






         """simple docstring"""


         A__										=  []
         for i in range(self.textLen - self.patLen + 1       ):
              A__										=  self.mismatch_in_text(UpperCamelCase       )
              if mismatch_index == -1:
                   positions.append(UpperCamelCase       )
              else:
                   A__										=  self.match_in_pattern(self.text[mismatch_index]       )
                   A__										=  (
                       mismatch_index - match_index
                   )  # shifting index lgtm [py/multiple-definition]
         return positions

SCREAMING_SNAKE_CASE_					:	List[Any]							  =						'ABAABA'
SCREAMING_SNAKE_CASE_					:	List[Any]							  =						'AB'
SCREAMING_SNAKE_CASE_					:	Union[str, Any]							  =						BoyerMooreSearch(text, pattern)
SCREAMING_SNAKE_CASE_					:	int							  =						bms.bad_character_heuristic()

if len(positions) == 0:
 print('No match found')
else:
 print('Pattern found in following positions: ')
 print(positions)






