




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

_A													=      [
    """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""",
]

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



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase				)						->    List[str]:
      lowerCAmelCase__      :			Union[str, Any]   =	{
          """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
      lowerCAmelCase__      :			Dict   =	int(re.match(R""".*layer_(\d*).*"""    ,		__UpperCAmelCase				)[1]				)
      layer_number -= 3
      return f"""h.{layer_number}.""" + key



def 	lowercase_				(			__UpperCAmelCase				)						->    List[str]:






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



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)						->    Optional[int]:






      if bloom_config_file == "":
            lowerCAmelCase__      :			Optional[int]   =	BloomConfig()
      else:
            lowerCAmelCase__      :			Union[str, Any]   =	BloomConfig.from_json_file(__UpperCAmelCase				)

      if shard_model:
            lowerCAmelCase__      :			Any   =	os.listdir(__UpperCAmelCase				)
            lowerCAmelCase__      :			Tuple   =	sorted(filter(lambda __UpperCAmelCase				: s.startswith("""layer"""				) and "model_00" in s    ,		__UpperCAmelCase				)				)

            lowerCAmelCase__      :			Dict   =	{"""weight_map""": {}, """metadata""": {}}
            lowerCAmelCase__      :			int   =	0

            lowerCAmelCase__      :			List[str]   =	None

            lowerCAmelCase__      :			Tuple   =	BloomConfig()

            for j, file in enumerate(__UpperCAmelCase				):
                  print("""Processing file: {}""".format(__UpperCAmelCase				)				)
                  lowerCAmelCase__      :			Union[str, Any]   =	None

                  for i in range(__UpperCAmelCase				):
                        # load all TP files
                        lowerCAmelCase__      :			int   =	file.replace("""model_00"""    ,		f"""model_0{i}"""				)
                        lowerCAmelCase__      :			Tuple   =	torch.load(os.path.join(__UpperCAmelCase    ,		__UpperCAmelCase				)    ,		map_location="""cpu"""				)

                        # Rename keys in the transformers names
                        lowerCAmelCase__      :			int   =	list(temp.keys()				)
                        for key in keys:
                              lowerCAmelCase__      :			Dict   =	temp.pop(__UpperCAmelCase				)

                        if tensors is None:
                              lowerCAmelCase__      :			List[str]   =	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
                                          lowerCAmelCase__      :			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
                                          lowerCAmelCase__      :			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				):
                              lowerCAmelCase__      :			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():
                        lowerCAmelCase__      :			Optional[int]   =	tensors[key]
                        total_size += value.numel() * get_dtype_size(value.dtype				)
                        if key not in index_dict["weight_map"]:
                              lowerCAmelCase__      :			List[str]   =	"""pytorch_model_{}-of-{}.bin""".format(
                                  str(j + 1				).zfill(5				)    ,		str(len(__UpperCAmelCase				)				).zfill(5				)				)

            lowerCAmelCase__      :			Dict   =	BloomConfig()
            lowerCAmelCase__      :			Any   =	pytorch_dump_folder_path + """/""" + CONFIG_NAME
            lowerCAmelCase__      :			Union[str, 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:
                  lowerCAmelCase__      :			Union[str, Any]   =	json.dumps(__UpperCAmelCase    ,		indent=2    ,		sort_keys=__UpperCAmelCase				) + """\n"""
                  f.write(__UpperCAmelCase				)
      else:
            lowerCAmelCase__      :			Optional[Any]   =	BloomModel(__UpperCAmelCase				)

            lowerCAmelCase__      :			Dict   =	os.listdir(__UpperCAmelCase				)
            lowerCAmelCase__      :			List[str]   =	sorted(filter(lambda __UpperCAmelCase				: s.startswith("""layer"""				) and "model_00" in s    ,		__UpperCAmelCase				)				)

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

                        # Rename keys in the transformers names
                        lowerCAmelCase__      :			Optional[Any]   =	list(temp.keys()				)
                        for key in keys:
                              lowerCAmelCase__      :			Optional[Any]   =	temp.pop(__UpperCAmelCase				)

                        if tensors is None:
                              lowerCAmelCase__      :			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
                                          lowerCAmelCase__      :			List[str]   =	1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN				) else 0
                                          # We concatenate these weights accross TP ranks
                                          lowerCAmelCase__      :			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				):
                              lowerCAmelCase__      :			Optional[int]   =	tensors[key] / pretraining_tp

                  lowerCAmelCase__      :			Tuple   =	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:
                        lowerCAmelCase__      :			int   =	set(other_keys.missing_keys				)
                  else:
                        lowerCAmelCase__      :			str   =	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				)
            lowerCAmelCase__      :			List[str]   =	pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
            lowerCAmelCase__      :			Optional[int]   =	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:
                  lowerCAmelCase__      :			Dict   =	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__":
   _A													=      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""",
   )
   _A													=      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 logging
import os
from typing import List, TextIO, Union

from conllu import parse_incr
from utils_ner import InputExample, Split, TokenClassificationTask


_A													=      logging.getLogger(__name__)


class _lowerCamelCase    (					a_						):






   def __init__(		self       :		List[Any]   ,     UpperCamelCase       :		Dict=-1   )			->						List[Any]:

         """simple docstring"""
         # in NER datasets, the last column is usually reserved for NER label
         lowerCAmelCase__      :			Optional[int]   =	label_idx






   def 	_lowerCAmelCase   (		self       :		Optional[int]   ,     UpperCamelCase       :		Optional[int]   ,     UpperCamelCase       :		Union[Split, str]   )			->						List[InputExample]:

         """simple docstring"""
         if isinstance(UpperCamelCase   ,     UpperCamelCase   ):
               lowerCAmelCase__      :			int   =	mode.value
         lowerCAmelCase__      :			Dict   =	os.path.join(UpperCamelCase   ,     f"""{mode}.txt"""   )
         lowerCAmelCase__      :			str   =	1
         lowerCAmelCase__      :			List[Any]   =	[]
         with open(UpperCamelCase   ,     encoding="""utf-8"""   ) as f:
               lowerCAmelCase__      :			str   =	[]
               lowerCAmelCase__      :			Any   =	[]
               for line in f:
                     if line.startswith("""-DOCSTART-"""   ) or line == "" or line == "\n":
                           if words:
                                 examples.append(InputExample(guid=f"""{mode}-{guid_index}"""   ,     words=UpperCamelCase   ,     labels=UpperCamelCase   )   )
                                 guid_index += 1
                                 lowerCAmelCase__      :			Optional[Any]   =	[]
                                 lowerCAmelCase__      :			Tuple   =	[]
                     else:
                           lowerCAmelCase__      :			Optional[int]   =	line.split(""" """   )
                           words.append(splits[0]   )
                           if len(UpperCamelCase   ) > 1:
                                 labels.append(splits[self.label_idx].replace("""\n"""   ,     """"""   )   )
                           else:
                                 # Examples could have no label for mode = "test"
                                 labels.append("""O"""   )
               if words:
                     examples.append(InputExample(guid=f"""{mode}-{guid_index}"""   ,     words=UpperCamelCase   ,     labels=UpperCamelCase   )   )
         return examples






   def 	_lowerCAmelCase   (		self       :		Any   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		List   )			->						List[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Dict   =	0
         for line in test_input_reader:
               if line.startswith("""-DOCSTART-"""   ) or line == "" or line == "\n":
                     writer.write(UpperCamelCase   )
                     if not preds_list[example_id]:
                           example_id += 1
               elif preds_list[example_id]:
                     lowerCAmelCase__      :			Union[str, Any]   =	line.split()[0] + """ """ + preds_list[example_id].pop(0   ) + """\n"""
                     writer.write(UpperCamelCase   )
               else:
                     logger.warning("""Maximum sequence length exceeded: No prediction for '%s'."""   ,     line.split()[0]   )






   def 	_lowerCAmelCase   (		self       :		str   ,     UpperCamelCase       :		str   )			->						List[str]:

         """simple docstring"""
         if path:
               with open(UpperCamelCase   ,     """r"""   ) as f:
                     lowerCAmelCase__      :			Any   =	f.read().splitlines()
               if "O" not in labels:
                     lowerCAmelCase__      :			List[str]   =	["""O"""] + labels
               return labels
         else:
               return ["O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]






class _lowerCamelCase    (					a_						):






   def __init__(		self       :		Union[str, Any]   )			->						Any:

         """simple docstring"""
         # in CONLL2003 dataset chunk column is second-to-last
         super().__init__(label_idx=-2   )






   def 	_lowerCAmelCase   (		self       :		Optional[Any]   ,     UpperCamelCase       :		str   )			->						List[str]:

         """simple docstring"""
         if path:
               with open(UpperCamelCase   ,     """r"""   ) as f:
                     lowerCAmelCase__      :			Any   =	f.read().splitlines()
               if "O" not in labels:
                     lowerCAmelCase__      :			str   =	["""O"""] + labels
               return labels
         else:
               return [
                   "O",
                   "B-ADVP",
                   "B-INTJ",
                   "B-LST",
                   "B-PRT",
                   "B-NP",
                   "B-SBAR",
                   "B-VP",
                   "B-ADJP",
                   "B-CONJP",
                   "B-PP",
                   "I-ADVP",
                   "I-INTJ",
                   "I-LST",
                   "I-PRT",
                   "I-NP",
                   "I-SBAR",
                   "I-VP",
                   "I-ADJP",
                   "I-CONJP",
                   "I-PP",
               ]






class _lowerCamelCase    (					a_						):






   def 	_lowerCAmelCase   (		self       :		Union[str, Any]   ,     UpperCamelCase       :		Any   ,     UpperCamelCase       :		Union[Split, str]   )			->						List[InputExample]:

         """simple docstring"""
         if isinstance(UpperCamelCase   ,     UpperCamelCase   ):
               lowerCAmelCase__      :			Optional[Any]   =	mode.value
         lowerCAmelCase__      :			int   =	os.path.join(UpperCamelCase   ,     f"""{mode}.txt"""   )
         lowerCAmelCase__      :			Optional[int]   =	1
         lowerCAmelCase__      :			List[Any]   =	[]

         with open(UpperCamelCase   ,     encoding="""utf-8"""   ) as f:
               for sentence in parse_incr(UpperCamelCase   ):
                     lowerCAmelCase__      :			Union[str, Any]   =	[]
                     lowerCAmelCase__      :			List[Any]   =	[]
                     for token in sentence:
                           words.append(token["""form"""]   )
                           labels.append(token["""upos"""]   )
                     assert len(UpperCamelCase   ) == len(UpperCamelCase   )
                     if words:
                           examples.append(InputExample(guid=f"""{mode}-{guid_index}"""   ,     words=UpperCamelCase   ,     labels=UpperCamelCase   )   )
                           guid_index += 1
         return examples






   def 	_lowerCAmelCase   (		self       :		List[str]   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		TextIO   ,     UpperCamelCase       :		List   )			->						List[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Any   =	0
         for sentence in parse_incr(UpperCamelCase   ):
               lowerCAmelCase__      :			Union[str, Any]   =	preds_list[example_id]
               lowerCAmelCase__      :			List[Any]   =	""""""
               for token in sentence:
                     out += f"""{token['form']} ({token['upos']}|{s_p.pop(0   )}) """
               out += "\n"
               writer.write(UpperCamelCase   )
               example_id += 1






   def 	_lowerCAmelCase   (		self       :		Dict   ,     UpperCamelCase       :		str   )			->						List[str]:

         """simple docstring"""
         if path:
               with open(UpperCamelCase   ,     """r"""   ) as f:
                     return f.read().splitlines()
         else:
               return [
                   "ADJ",
                   "ADP",
                   "ADV",
                   "AUX",
                   "CCONJ",
                   "DET",
                   "INTJ",
                   "NOUN",
                   "NUM",
                   "PART",
                   "PRON",
                   "PROPN",
                   "PUNCT",
                   "SCONJ",
                   "SYM",
                   "VERB",
                   "X",
               ]





