

"""simple docstring"""






import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

__lowerCamelCase										=		[
    "word_embeddings_layernorm.weight",
    "word_embeddings_layernorm.bias",
    "input_layernorm.weight",
    "input_layernorm.bias",
    "post_attention_layernorm.weight",
    "post_attention_layernorm.bias",
    "self_attention.dense.bias",
    "mlp.dense_4h_to_h.bias",
    "ln_f.weight",
    "ln_f.bias",
]

__lowerCamelCase										=		[
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]





def 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase						) ->      List[str]:
 __magic_name__     =      {
     '''word_embeddings.weight''': '''word_embeddings.weight''',
     '''word_embeddings.norm.weight''': '''word_embeddings_layernorm.weight''',
     '''word_embeddings.norm.bias''': '''word_embeddings_layernorm.bias''',
     '''weight''': '''ln_f.weight''',
     '''bias''': '''ln_f.bias''',
 }

 if key in layer_rename_map:
  return layer_rename_map[key]

 # Handle transformer blocks
 __magic_name__     =      int(re.match(r'''.*layer_(\d*).*'''      ,   a_						)[1]						)
 layer_number -= 3
 return f'''h.{layer_number}.''' + key





def 	lowercase      (   __UpperCamelCase						) ->      List[str]:
 if dtype == torch.bool:
  return 1 / 8
 __magic_name__     =      re.search(r'''[^\d](\d+)$'''      ,   str(a_						)						)
 if bit_search is None:
  raise ValueError(f'''`dtype` is not a valid dtype: {dtype}.'''						)
 __magic_name__     =      int(bit_search.groups()[0]						)
 return bit_size // 8





def 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase      ,   __UpperCamelCase      ,   __UpperCamelCase      ,   __UpperCamelCase						) ->      int:
 # Construct model
 if bloom_config_file == "":
  __magic_name__     =      BloomConfig()
 else:
  __magic_name__     =      BloomConfig.from_json_file(a_						)

 if shard_model:
  __magic_name__     =      os.listdir(a_						)
  __magic_name__     =      sorted(filter(lambda __UpperCamelCase						: s.startswith('''layer'''						) and "model_00" in s      ,   a_						)						)

  __magic_name__     =      {'''weight_map''': {}, '''metadata''': {}}
  __magic_name__     =      0

  __magic_name__     =      None

  __magic_name__     =      BloomConfig()

  for j, file in enumerate(a_						):
   print('''Processing file: {}'''.format(a_						)						)
   __magic_name__     =      None

   for i in range(a_						):
    # load all TP files
    __magic_name__     =      file.replace('''model_00'''      ,   f'''model_0{i}'''						)
    __magic_name__     =      torch.load(os.path.join(a_      ,   a_						)      ,   map_location='''cpu'''						)

    # Rename keys in the transformers names
    __magic_name__     =      list(temp.keys()						)
    for key in keys:
     __magic_name__     =      temp.pop(a_						)

    if tensors is None:
     __magic_name__     =      temp
    else:
     for key in tensors.keys():
      if any(key.endswith(a_						) for end in WEIGHTS_TO_AVERAGE_ENDSWITH						):
       # We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
       tensors[key] += temp[key]
      else:
       # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
       __magic_name__     =      1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN						) else 0
       # We concatenate these weights accross TP ranks
       __magic_name__     =      torch.cat([tensors[key], temp[key]]      ,   dim=a_						)

            # Divide by the number of TP the weights we want to average
   for key in tensors.keys():
    if any(key.endswith(a_						) for end in WEIGHTS_TO_AVERAGE_ENDSWITH						):
     __magic_name__     =      tensors[key] / pretraining_tp
   torch.save(
       a_      ,   os.path.join(
           a_      ,   '''pytorch_model_{}-of-{}.bin'''.format(str(j + 1						).zfill(5						)      ,   str(len(a_						)						).zfill(5						)						)      ,   )      ,   )

   for key in tensors.keys():
    __magic_name__     =      tensors[key]
    total_size += value.numel() * get_dtype_size(value.dtype						)
    if key not in index_dict["weight_map"]:
     __magic_name__     =      '''pytorch_model_{}-of-{}.bin'''.format(
         str(j + 1						).zfill(5						)      ,   str(len(a_						)						).zfill(5						)						)

  __magic_name__     =      BloomConfig()
  __magic_name__     =      pytorch_dump_folder_path + '''/''' + CONFIG_NAME
  __magic_name__     =      total_size
  with open(a_      ,   '''w'''      ,   encoding='''utf-8'''						) as f:
   f.write(config.to_json_string()						)
  with open(os.path.join(a_      ,   WEIGHTS_NAME + '''.index.json'''						)      ,   '''w'''      ,   encoding='''utf-8'''						) as f:
   __magic_name__     =      json.dumps(a_      ,   indent=2      ,   sort_keys=a_						) + '''\n'''
   f.write(a_						)
 else:
  __magic_name__     =      BloomModel(a_						)

  __magic_name__     =      os.listdir(a_						)
  __magic_name__     =      sorted(filter(lambda __UpperCamelCase						: s.startswith('''layer'''						) and "model_00" in s      ,   a_						)						)

  __magic_name__     =      None
  for i, file in enumerate(a_						):
   __magic_name__     =      None
   for i in range(a_						):
    # load all TP files
    __magic_name__     =      file.replace('''model_00'''      ,   f'''model_0{i}'''						)
    __magic_name__     =      torch.load(os.path.join(a_      ,   a_						)      ,   map_location='''cpu'''						)

    # Rename keys in the transformers names
    __magic_name__     =      list(temp.keys()						)
    for key in keys:
     __magic_name__     =      temp.pop(a_						)

    if tensors is None:
     __magic_name__     =      temp
    else:
     for key in tensors.keys():
      # We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
      if any(key.endswith(a_						) for end in WEIGHTS_TO_AVERAGE_ENDSWITH						):
       tensors[key] += temp[key]
      else:
       # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
       __magic_name__     =      1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN						) else 0
       # We concatenate these weights accross TP ranks
       __magic_name__     =      torch.cat([tensors[key], temp[key]]      ,   dim=a_						)

            # Divide by the number of TP the weights we want to average
   for key in tensors.keys():
    if any(key.endswith(a_						) for end in WEIGHTS_TO_AVERAGE_ENDSWITH						):
     __magic_name__     =      tensors[key] / pretraining_tp

   __magic_name__     =      model.load_state_dict(a_      ,   strict=a_						)
   assert not other_keys.unexpected_keys, f'''The keys {other_keys.unexpected_keys} are unexpected'''
   if missing_keys is None:
    __magic_name__     =      set(other_keys.missing_keys						)
   else:
    __magic_name__     =      missing_keys.intersection(set(other_keys.missing_keys						)						)

  assert not missing_keys, f'''The keys {missing_keys} are missing'''

  # Save pytorch-model
  os.makedirs(a_      ,   exist_ok=a_						)
  __magic_name__     =      pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
  __magic_name__     =      pytorch_dump_folder_path + '''/''' + CONFIG_NAME
  print(f'''Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}'''						)
  if config.torch_dtype is not None:
   __magic_name__     =      model.to(config.torch_dtype						)
  torch.save(model.state_dict()      ,   a_						)
  print(f'''Save configuration file to {pytorch_config_dump_path}'''						)
  with open(a_      ,   '''w'''      ,   encoding='''utf-8'''						) as f:
   f.write(config.to_json_string()						)


if __name__ == "__main__":
 __lowerCamelCase										=		argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--bloom_checkpoint_path",
     default=None,
     type=str,
     required=True,
     help="Path to the Megatron-LM checkpoint path.",
 )
 parser.add_argument(
     "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
 )
 parser.add_argument(
     "--bloom_config_file",
     default="",
     type=str,
     help=(
         "An optional config json file corresponding to the pre-trained model. \n"
         "This specifies the model architecture."
     ),
 )
 parser.add_argument(
     "--shard_model",
     action="store_true",
     help="An optional setting to shard the output model \nThis enables sharding the converted checkpoint",
 )
 parser.add_argument(
     "--pretraining_tp",
     default=4,
     type=int,
     help="Pretraining TP rank that has been used when training the model in Megatron-LM \n",
 )
 __lowerCamelCase										=		parser.parse_args()
 convert_bloom_checkpoint_to_pytorch(
     args.bloom_checkpoint_path,
     args.bloom_config_file,
     args.pytorch_dump_folder_path,
     args.shard_model,
     args.pretraining_tp,
 )




"""simple docstring"""






import os
import unittest

from transformers import FunnelTokenizer, FunnelTokenizerFast
from transformers.models.funnel.tokenization_funnel import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin
@require_tokenizers
class        _lowercase					(     __UpperCAmelCase   , unittest.TestCase     ):
     _lowerCamelCase     =       FunnelTokenizer
     _lowerCamelCase     =       FunnelTokenizerFast
     _lowerCamelCase     =       True
     _lowerCamelCase     =       True







     def 	lowerCAmelCase__    (      self       ):
      super().setUp()

      __magic_name__     =      [
          '''<unk>''',
          '''<cls>''',
          '''<sep>''',
          '''want''',
          '''##want''',
          '''##ed''',
          '''wa''',
          '''un''',
          '''runn''',
          '''##ing''',
          ''',''',
          '''low''',
          '''lowest''',
      ]
      __magic_name__     =      os.path.join(self.tmpdirname ,						VOCAB_FILES_NAMES['''vocab_file''']       )
      with open(self.vocab_file ,						'''w''' ,						encoding='''utf-8'''       ) as vocab_writer:
       vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]       )       )







     def 	lowerCAmelCase__    (      self ,						**UpperCamelCase_       ):
      return FunnelTokenizer.from_pretrained(self.tmpdirname ,						**UpperCamelCase_       )







     def 	lowerCAmelCase__    (      self ,						**UpperCamelCase_       ):
      return FunnelTokenizerFast.from_pretrained(self.tmpdirname ,						**UpperCamelCase_       )







     def 	lowerCAmelCase__    (      self ,						UpperCamelCase_       ):
      __magic_name__     =      '''UNwant\u00E9d,running'''
      __magic_name__     =      '''unwanted, running'''
      return input_text, output_text







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.tokenizer_class(self.vocab_file       )

      __magic_name__     =      tokenizer.tokenize('''UNwant\u00E9d,running'''       )
      self.assertListEqual(UpperCamelCase_ ,						['''un''', '''##want''', '''##ed''', ''',''', '''runn''', '''##ing''']       )
      self.assertListEqual(tokenizer.convert_tokens_to_ids(UpperCamelCase_       ) ,						[7, 4, 5, 10, 8, 9]       )







     def 	lowerCAmelCase__    (      self       ):
      __magic_name__     =      self.get_tokenizers(do_lower_case=UpperCamelCase_       )
      for tokenizer in tokenizers:
       __magic_name__     =      tokenizer('''UNwant\u00E9d,running'''       )
       __magic_name__     =      len(inputs['''input_ids''']       ) - 1
       self.assertListEqual(inputs['''token_type_ids'''] ,						[2] + [0] * sentence_len       )

       __magic_name__     =      tokenizer('''UNwant\u00E9d,running''' ,						'''UNwant\u00E9d,running'''       )
       self.assertListEqual(inputs['''token_type_ids'''] ,						[2] + [0] * sentence_len + [1] * sentence_len       )

