




"""simple docstring"""


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

import torch

from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer


try:
     from transformers import LlamaTokenizerFast
except ImportError as e:
     warnings.warn(e)
     warnings.warn(
         '''The converted tokenizer will be the `slow` tokenizer. To use the fast, update your `tokenizers` library and re-run the tokenizer conversion'''
     )
     A_      =    None
A_      =    {
    '''7B''': 1_10_08,
    '''13B''': 1_38_24,
    '''30B''': 1_79_20,
    '''65B''': 2_20_16,
    '''70B''': 2_86_72,
}
A_      =    {
    '''7B''': 1,
    '''7Bf''': 1,
    '''13B''': 2,
    '''13Bf''': 2,
    '''30B''': 4,
    '''65B''': 8,
    '''70B''': 8,
    '''70Bf''': 8,
}


def  UpperCAmelCase__	(snake_case__							: int		,       snake_case__							: List[Any]=1		,       snake_case__							: List[Any]=2_56   ):
  """simple docstring"""

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


def  UpperCAmelCase__	(snake_case__							: str   ):
  """simple docstring"""

  with open(a_		,       """r"""   ) as f:
    return json.load(a_   )


def  UpperCAmelCase__	(snake_case__							: Dict		,       snake_case__							: Tuple   ):
  """simple docstring"""

  with open(a_		,       """w"""   ) as f:
    json.dump(a_		,       a_   )


def  UpperCAmelCase__	(snake_case__							: List[str]		,       snake_case__							: Union[str, Any]		,       snake_case__							: Union[str, Any]		,       snake_case__							: Optional[int]=True   ):
  """simple docstring"""

  os.makedirs(a_		,       exist_ok=a_   )
  _snake_case							:		Optional[Any]								=  os.path.join(a_		,       """tmp"""   )
  os.makedirs(a_		,       exist_ok=a_   )

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

  if "n_kv_heads" in params:
    _snake_case							:		Any								=  params["""n_kv_heads"""]  # for GQA / MQA
    _snake_case							:		str								=  n_heads_per_shard // num_key_value_heads
    _snake_case							:		Union[str, Any]								=  dim // num_key_value_heads
  else:  # compatibility with other checkpoints
    _snake_case							:		Tuple								=  n_heads
    _snake_case							:		List[str]								=  n_heads_per_shard
    _snake_case							:		Tuple								=  dim

  # permute for sliced rotary
  def permute(snake_case__							: Union[str, Any]		,       snake_case__							: Union[str, Any]=n_heads		,       snake_case__							: Union[str, Any]=dim		,       snake_case__							: int=dim   ):
    return w.view(a_		,       dima // n_heads // 2		,       2		,       a_   ).transpose(1		,       2   ).reshape(a_		,       a_   )

  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							:		Dict								=  torch.load(os.path.join(a_		,       """consolidated.00.pth"""   )		,       map_location="""cpu"""   )
  else:
    # Sharded
    _snake_case							:		List[Any]								=  [
        torch.load(os.path.join(a_		,       F"consolidated.{i:02d}.pth"   )		,       map_location="""cpu"""   )
        for i in range(a_   )
    ]
  _snake_case							:		Optional[int]								=  0
  _snake_case							:		Optional[Any]								=  {"""weight_map""": {}}
  for layer_i in range(a_   ):
    _snake_case							:		Dict								=  F"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
    if model_size == "7B":
      # Unsharded
      _snake_case							:		Union[str, Any]								=  {
          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							:		Optional[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							:		Dict								=  permute(
          torch.cat(
              [
                  loaded[i][F"layers.{layer_i}.attention.wq.weight"].view(a_		,       a_		,       a_   )
                  for i in range(a_   )
              ]		,       dim=0		,       ).reshape(a_		,       a_   )   )
      _snake_case							:		str								=  permute(
          torch.cat(
              [
                  loaded[i][F"layers.{layer_i}.attention.wk.weight"].view(
                      a_		,       a_		,       a_   )
                  for i in range(a_   )
              ]		,       dim=0		,       ).reshape(a_		,       a_   )		,       a_		,       a_		,       a_		,       )
      _snake_case							:		Optional[Any]								=  torch.cat(
          [
              loaded[i][F"layers.{layer_i}.attention.wv.weight"].view(
                  a_		,       a_		,       a_   )
              for i in range(a_   )
          ]		,       dim=0		,       ).reshape(a_		,       a_   )

      _snake_case							:		Union[str, Any]								=  torch.cat(
          [loaded[i][F"layers.{layer_i}.attention.wo.weight"] for i in range(a_   )]		,       dim=1   )
      _snake_case							:		Any								=  torch.cat(
          [loaded[i][F"layers.{layer_i}.feed_forward.w1.weight"] for i in range(a_   )]		,       dim=0   )
      _snake_case							:		int								=  torch.cat(
          [loaded[i][F"layers.{layer_i}.feed_forward.w2.weight"] for i in range(a_   )]		,       dim=1   )
      _snake_case							:		str								=  torch.cat(
          [loaded[i][F"layers.{layer_i}.feed_forward.w3.weight"] for i in range(a_   )]		,       dim=0   )

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

  _snake_case							:		Union[str, Any]								=  F"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
  if model_size == "7B":
    # Unsharded
    _snake_case							:		int								=  {
        """model.embed_tokens.weight""": loaded["""tok_embeddings.weight"""],
        """model.norm.weight""": loaded["""norm.weight"""],
        """lm_head.weight""": loaded["""output.weight"""],
    }
  else:
    _snake_case							:		Optional[Any]								=  {
        """model.norm.weight""": loaded[0]["""norm.weight"""],
        """model.embed_tokens.weight""": torch.cat(
            [loaded[i]["""tok_embeddings.weight"""] for i in range(a_   )]		,       dim=1   ),
        """lm_head.weight""": torch.cat([loaded[i]["""output.weight"""] for i in range(a_   )]		,       dim=0   ),
    }

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

  # Write configs
  _snake_case							:		Union[str, Any]								=  {"""total_size""": param_count * 2}
  write_json(a_		,       os.path.join(a_		,       """pytorch_model.bin.index.json"""   )   )
  _snake_case							:		Tuple								=  params["""ffn_dim_multiplier"""] if """ffn_dim_multiplier""" in params else 1
  _snake_case							:		int								=  params["""multiple_of"""] if """multiple_of""" in params else 2_56
  _snake_case							:		str								=  LlamaConfig(
      hidden_size=a_		,       intermediate_size=compute_intermediate_size(a_		,       a_		,       a_   )		,       num_attention_heads=params["""n_heads"""]		,       num_hidden_layers=params["""n_layers"""]		,       rms_norm_eps=params["""norm_eps"""]		,       num_key_value_heads=a_		,       )
  config.save_pretrained(a_   )

  # 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							:		str								=  LlamaForCausalLM.from_pretrained(a_		,       torch_dtype=torch.floataa		,       low_cpu_mem_usage=a_   )
  # Avoid saving this as part of the config.
  del model.config._name_or_path

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


def  UpperCAmelCase__	(snake_case__							: Tuple		,       snake_case__							: Optional[Any]   ):
  """simple docstring"""

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


def  UpperCAmelCase__	():
  """simple docstring"""

  _snake_case							:		Tuple								=  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=a_		,       help="""Whether or not to save using `safetensors`."""   )
  _snake_case							:		Union[str, 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		,       )
  _snake_case							:		Optional[int]								=  os.path.join(args.input_dir		,       """tokenizer.model"""   )
  write_tokenizer(args.output_dir		,       a_   )


if __name__ == "__main__":
     main()









"""simple docstring"""


import os
import re
import shutil
from argparse import ArgumentParser, Namespace

from datasets.commands import BaseDatasetsCLICommand
from datasets.utils.logging import get_logger


A_      =    '''<<<<<<< This should probably be modified because it mentions: '''

A_      =    '''=======
>>>>>>>
'''

A_      =    [
    '''TextEncoderConfig''',
    '''ByteTextEncoder''',
    '''SubwordTextEncoder''',
    '''encoder_config''',
    '''maybe_build_from_corpus''',
    '''manual_dir''',
]

A_      =    [
    # (pattern, replacement)
    # Order is important here for some replacements
    (r'''tfds\.core''', r'''datasets'''),
    (r'''tf\.io\.gfile\.GFile''', r'''open'''),
    (r'''tf\.([\w\d]+)''', r'''datasets.Value(\'\1\')'''),
    (r'''tfds\.features\.Text\(\)''', r'''datasets.Value(\'string\')'''),
    (r'''tfds\.features\.Text\(''', r'''datasets.Value(\'string\'),'''),
    (r'''features\s*=\s*tfds.features.FeaturesDict\(''', r'''features=datasets.Features('''),
    (r'''tfds\.features\.FeaturesDict\(''', r'''dict('''),
    (r'''The TensorFlow Datasets Authors''', r'''The TensorFlow Datasets Authors and the HuggingFace Datasets Authors'''),
    (r'''tfds\.''', r'''datasets.'''),
    (r'''dl_manager\.manual_dir''', r'''self.config.data_dir'''),
    (r'''self\.builder_config''', r'''self.config'''),
]


def  UpperCAmelCase__	(snake_case__							: Namespace   ):
  """simple docstring"""

  return ConvertCommand(args.tfds_path		,       args.datasets_directory   )






class  lowercase(					__a			):







 '''simple docstring'''





 @staticmethod
 def 			UpperCamelCase_       (    a_:		ArgumentParser		):




   '''simple docstring'''



   _snake_case							:		Tuple								=  parser.add_parser(
       """convert""",				help="""Convert a TensorFlow Datasets dataset to a HuggingFace Datasets dataset.""",				)
   train_parser.add_argument(
       """--tfds_path""",				type=a_,				required=a_,				help="""Path to a TensorFlow Datasets folder to convert or a single tfds file to convert.""",				)
   train_parser.add_argument(
       """--datasets_directory""",				type=a_,				required=a_,				help="""Path to the HuggingFace Datasets folder."""		)
   train_parser.set_defaults(func=a_		)


 def __init__(    self:		List[str],				a_:		str,				a_:		str,				*a_:		str		):




   '''simple docstring'''



   _snake_case							:		Optional[Any]								=  get_logger("""datasets-cli/converting"""		)

   _snake_case							:		Any								=  tfds_path
   _snake_case							:		Optional[Any]								=  datasets_directory


 def 			UpperCamelCase_       (    self:		Optional[int]		):




   '''simple docstring'''



   if os.path.isdir(self._tfds_path		):
     _snake_case							:		int								=  os.path.abspath(self._tfds_path		)
   elif os.path.isfile(self._tfds_path		):
     _snake_case							:		Any								=  os.path.dirname(self._tfds_path		)
   else:
     raise ValueError("""--tfds_path is neither a directory nor a file. Please check path."""		)

   _snake_case							:		Union[str, Any]								=  os.path.abspath(self._datasets_directory		)

   self._logger.info(f"Converting datasets from {abs_tfds_path} to {abs_datasets_path}"		)

   _snake_case							:		Tuple								=  []
   _snake_case							:		Dict								=  []
   _snake_case							:		Optional[Any]								=  {}

   if os.path.isdir(self._tfds_path		):
     _snake_case							:		List[str]								=  os.listdir(a_		)
   else:
     _snake_case							:		int								=  [os.path.basename(self._tfds_path		)]

   for f_name in file_names:
     self._logger.info(f"Looking at file {f_name}"		)
     _snake_case							:		Dict								=  os.path.join(a_,				a_		)
     _snake_case							:		Union[str, Any]								=  os.path.join(a_,				a_		)

     if not os.path.isfile(a_		) or "__init__" in f_name or "_test" in f_name or ".py" not in f_name:
       self._logger.info("""Skipping file"""		)
       continue

     with open(a_,				encoding="""utf-8"""		) as f:
       _snake_case							:		str								=  f.readlines()

     _snake_case							:		List[str]								=  []
     _snake_case							:		Any								=  False
     _snake_case							:		Union[str, Any]								=  False
     _snake_case							:		Optional[Any]								=  []
     for line in lines:
       _snake_case							:		Optional[int]								=  line

       # Convert imports
       if "import tensorflow.compat.v2 as tf" in out_line:
         continue
       elif "@tfds.core" in out_line:
         continue
       elif "builder=self" in out_line:
         continue
       elif "import tensorflow_datasets.public_api as tfds" in out_line:
         _snake_case							:		Optional[Any]								=  """import datasets\n"""
       elif "import tensorflow" in out_line:
         # order is important here
         _snake_case							:		Optional[int]								=  """"""
         continue
       elif "from absl import logging" in out_line:
         _snake_case							:		int								=  """from datasets import logging\n"""
       elif "getLogger" in out_line:
         _snake_case							:		Any								=  out_line.replace("""getLogger""",				"""get_logger"""		)
       elif any(expression in out_line for expression in TO_HIGHLIGHT		):
         _snake_case							:		Union[str, Any]								=  True
         _snake_case							:		Optional[Any]								=  list(filter(lambda a_		: e in out_line,				a_		)		)
         out_lines.append(HIGHLIGHT_MESSAGE_PRE + str(a_		) + """\n"""		)
         out_lines.append(a_		)
         out_lines.append(a_		)
         continue
       else:
         for pattern, replacement in TO_CONVERT:
           _snake_case							:		List[str]								=  re.sub(a_,				a_,				a_		)

                # Take care of saving utilities (to later move them together with main script)
       if "tensorflow_datasets" in out_line:
         _snake_case							:		Dict								=  re.match(r"""from\stensorflow_datasets.*import\s([^\.\r\n]+)""",				a_		)
         tfds_imports.extend(imp.strip() for imp in match.group(1		).split(""","""		)		)
         _snake_case							:		Optional[Any]								=  """from . import """ + match.group(1		)

       # Check we have not forget anything
       if "tf." in out_line or "tfds." in out_line or "tensorflow_datasets" in out_line:
         raise ValueError(f"Error converting {out_line.strip()}"		)

       if "GeneratorBasedBuilder" in out_line or "BeamBasedBuilder" in out_line:
         _snake_case							:		Tuple								=  True
       out_lines.append(a_		)

     if is_builder or "wmt" in f_name:
       # We create a new directory for each dataset
       _snake_case							:		List[str]								=  f_name.replace(""".py""",				""""""		)
       _snake_case							:		str								=  os.path.join(a_,				a_		)
       _snake_case							:		str								=  os.path.join(a_,				a_		)
       os.makedirs(a_,				exist_ok=a_		)
       self._logger.info(f"Adding directory {output_dir}"		)
       imports_to_builder_map.update({imp: output_dir for imp in tfds_imports}		)
     else:
       # Utilities will be moved at the end
       utils_files.append(a_		)

     if needs_manual_update:
       with_manual_update.append(a_		)

     with open(a_,				"""w""",				encoding="""utf-8"""		) as f:
       f.writelines(a_		)
     self._logger.info(f"Converted in {output_file}"		)

   for utils_file in utils_files:
     try:
       _snake_case							:		Optional[int]								=  os.path.basename(a_		)
       _snake_case							:		Optional[Any]								=  imports_to_builder_map[f_name.replace(""".py""",				""""""		)]
       self._logger.info(f"Moving {dest_folder} to {utils_file}"		)
       shutil.copy(a_,				a_		)
     except KeyError:
       self._logger.error(f"Cannot find destination folder for {utils_file}. Please copy manually."		)

   if with_manual_update:
     for file_path in with_manual_update:
       self._logger.warning(
           f"You need to manually update file {file_path} to remove configurations using 'TextEncoderConfig'."		)


