


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

UpperCAmelCase_			:  List[Any]										=   [
    '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',
]

UpperCAmelCase_			:  str										=   [
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]


def       _lowercase							(					UpperCamelCase__							:       Any,					UpperCamelCase__							:       Optional[int]      ):
 __A	:		int	  =      {
     '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
 __A	:		Optional[Any]	  =      int(re.match(r'.*layer_(\d*).*',					UpperCamelCase__      )[1]      )
 layer_number -= 3
 return f"""h.{layer_number}.""" + key


def       _lowercase							(					UpperCamelCase__							:       Dict      ):
 if dtype == torch.bool:
  return 1 / 8
 __A	:		Optional[Any]	  =      re.search(r'[^\d](\d+)$',					str(UpperCamelCase__      )      )
 if bit_search is None:
  raise ValueError(f"""`dtype` is not a valid dtype: {dtype}."""      )
 __A	:		Any	  =      int(bit_search.groups()[0]      )
 return bit_size // 8


def       _lowercase							(					UpperCamelCase__							:       Any,					UpperCamelCase__							:       Union[str, Any],					UpperCamelCase__							:       Tuple,					UpperCamelCase__							:       int,					UpperCamelCase__							:       str      ):
 # Construct model
 if bloom_config_file == "":
  __A	:		List[str]	  =      BloomConfig()
 else:
  __A	:		Optional[int]	  =      BloomConfig.from_json_file(UpperCamelCase__      )

 if shard_model:
  __A	:		Any	  =      os.listdir(UpperCamelCase__      )
  __A	:		int	  =      sorted(filter(lambda UpperCamelCase__      : s.startswith('layer'      ) and "model_00" in s,					UpperCamelCase__      )      )

  __A	:		Optional[Any]	  =      {'weight_map': {}, 'metadata': {}}
  __A	:		Optional[Any]	  =      0

  __A	:		List[Any]	  =      None

  __A	:		str	  =      BloomConfig()

  for j, file in enumerate(UpperCamelCase__      ):
   print('Processing file: {}'.format(UpperCamelCase__      )      )
   __A	:		Dict	  =      None

   for i in range(UpperCamelCase__      ):
    # load all TP files
    __A	:		Any	  =      file.replace('model_00',					f"""model_0{i}"""      )
    __A	:		List[Any]	  =      torch.load(os.path.join(UpperCamelCase__,					UpperCamelCase__      ),					map_location='cpu'      )

    # Rename keys in the transformers names
    __A	:		Tuple	  =      list(temp.keys()      )
    for key in keys:
     __A	:		Any	  =      temp.pop(UpperCamelCase__      )

    if tensors is None:
     __A	:		Tuple	  =      temp
    else:
     for key in tensors.keys():
      if any(key.endswith(UpperCamelCase__      ) 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
       __A	:		Union[str, Any]	  =      1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN      ) else 0
       # We concatenate these weights accross TP ranks
       __A	:		List[Any]	  =      torch.cat([tensors[key], temp[key]],					dim=UpperCamelCase__      )

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

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

  __A	:		List[Any]	  =      BloomConfig()
  __A	:		Optional[Any]	  =      pytorch_dump_folder_path + '/' + CONFIG_NAME
  __A	:		Optional[Any]	  =      total_size
  with open(UpperCamelCase__,					'w',					encoding='utf-8'      ) as f:
   f.write(config.to_json_string()      )
  with open(os.path.join(UpperCamelCase__,					WEIGHTS_NAME + '.index.json'      ),					'w',					encoding='utf-8'      ) as f:
   __A	:		List[str]	  =      json.dumps(UpperCamelCase__,					indent=2,					sort_keys=UpperCamelCase__      ) + '\n'
   f.write(UpperCamelCase__      )
 else:
  __A	:		int	  =      BloomModel(UpperCamelCase__      )

  __A	:		Tuple	  =      os.listdir(UpperCamelCase__      )
  __A	:		Any	  =      sorted(filter(lambda UpperCamelCase__      : s.startswith('layer'      ) and "model_00" in s,					UpperCamelCase__      )      )

  __A	:		Dict	  =      None
  for i, file in enumerate(UpperCamelCase__      ):
   __A	:		Union[str, Any]	  =      None
   for i in range(UpperCamelCase__      ):
    # load all TP files
    __A	:		Dict	  =      file.replace('model_00',					f"""model_0{i}"""      )
    __A	:		str	  =      torch.load(os.path.join(UpperCamelCase__,					UpperCamelCase__      ),					map_location='cpu'      )

    # Rename keys in the transformers names
    __A	:		int	  =      list(temp.keys()      )
    for key in keys:
     __A	:		Union[str, Any]	  =      temp.pop(UpperCamelCase__      )

    if tensors is None:
     __A	:		Union[str, Any]	  =      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(UpperCamelCase__      ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH      ):
       tensors[key] += temp[key]
      else:
       # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
       __A	:		Optional[int]	  =      1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN      ) else 0
       # We concatenate these weights accross TP ranks
       __A	:		int	  =      torch.cat([tensors[key], temp[key]],					dim=UpperCamelCase__      )

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

   __A	:		str	  =      model.load_state_dict(UpperCamelCase__,					strict=UpperCamelCase__      )
   assert not other_keys.unexpected_keys, f"""The keys {other_keys.unexpected_keys} are unexpected"""
   if missing_keys is None:
    __A	:		Optional[int]	  =      set(other_keys.missing_keys      )
   else:
    __A	:		List[Any]	  =      missing_keys.intersection(set(other_keys.missing_keys      )      )

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

  # Save pytorch-model
  os.makedirs(UpperCamelCase__,					exist_ok=UpperCamelCase__      )
  __A	:		Optional[Any]	  =      pytorch_dump_folder_path + '/' + WEIGHTS_NAME
  __A	:		Dict	  =      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:
   __A	:		Tuple	  =      model.to(config.torch_dtype      )
  torch.save(model.state_dict(),					UpperCamelCase__      )
  print(f"""Save configuration file to {pytorch_config_dump_path}"""      )
  with open(UpperCamelCase__,					'w',					encoding='utf-8'      ) as f:
   f.write(config.to_json_string()      )


if __name__ == "__main__":
 UpperCAmelCase_			:  Optional[Any]										=   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',
 )
 UpperCAmelCase_			:  Tuple										=   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'''
def       _lowercase							(					UpperCamelCase__							:       dict      ):
 __A	:		Dict	  =      set()
 # edges = list of graph's edges
 __A	:		Any	  =      get_edges(UpperCamelCase__      )

 # While there are still elements in edges list, take an arbitrary edge
 # (from_node, to_node) and add his extremity to chosen_vertices and then
 # remove all arcs adjacent to the from_node and to_node
 while edges:
  __A  ,__A	:		List[Any]	  =      edges.pop()
  chosen_vertices.add(UpperCamelCase__      )
  chosen_vertices.add(UpperCamelCase__      )
  for edge in edges.copy():
   if from_node in edge or to_node in edge:
    edges.discard(UpperCamelCase__      )
 return chosen_vertices


def       _lowercase							(					UpperCamelCase__							:       dict      ):
 __A	:		Any	  =      set()
 for from_node, to_nodes in graph.items():
  for to_node in to_nodes:
   edges.add((from_node, to_node)      )
 return edges


if __name__ == "__main__":
 import doctest

 doctest.testmod()

 # graph = {0: [1, 3], 1: [0, 3], 2: [0, 3, 4], 3: [0, 1, 2], 4: [2, 3]}
 # print(f"Matching vertex cover:\n{matching_min_vertex_cover(graph)}")






