




import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only
from utils_rag import save_json







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


     UpperCAmelCase_     :				Optional[Any]  =     filter(lambda _lowercase : p.requires_grad					,		model.parameters() )
     UpperCAmelCase_     :				Dict  =     sum([np.prod(p.size() ) for p in model_parameters] )
     return params


__a							     =					logging.getLogger(__name__)







def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     if metric == "rouge2":
          UpperCAmelCase_     :				Any  =     '''{val_avg_rouge2:.4f}-{step_count}'''
     elif metric == "bleu":
          UpperCAmelCase_     :				Union[str, Any]  =     '''{val_avg_bleu:.4f}-{step_count}'''
     elif metric == "em":
          UpperCAmelCase_     :				Dict  =     '''{val_avg_em:.4f}-{step_count}'''
     elif metric == "loss":
          UpperCAmelCase_     :				Optional[int]  =     '''{val_avg_loss:.4f}-{step_count}'''
     else:
          raise NotImplementedError(
              f'''seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this'''
              ''' function.''' )

     UpperCAmelCase_     :				Union[str, Any]  =     ModelCheckpoint(
         dirpath=_lowercase					,		filename=_lowercase					,		monitor=f'''val_{metric}'''					,		mode='''max'''					,		save_top_k=1					,		every_n_epochs=1					,		)

     return checkpoint_callback







def lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


     return EarlyStopping(
         monitor=f'''val_{metric}'''					,		mode='''min''' if '''loss''' in metric else '''max'''					,		patience=_lowercase					,		verbose=_lowercase					,		)



class     __a(	pl.Callback  ):






    """simple docstring"""





    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )					->						Any:
         UpperCAmelCase_     :				Any  =     {f'''lr_group_{i}''': param['''lr'''] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups    )}
         pl_module.logger.log_metrics(_SCREAMING_SNAKE_CASE    )





    @rank_zero_only
    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE=True    )					->						None:
         logger.info(f'''***** {type_path} results at step {trainer.global_step:05d} *****'''    )
         UpperCAmelCase_     :				str  =     trainer.callback_metrics
         trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['''log''', '''progress_bar''', '''preds''']}    )
         # Log results
         UpperCAmelCase_     :				Tuple  =     Path(pl_module.hparams.output_dir    )
         if type_path == "test":
              UpperCAmelCase_     :				str  =     od / '''test_results.txt'''
              UpperCAmelCase_     :				Dict  =     od / '''test_generations.txt'''
         else:
              # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
              # If people want this it will be easy enough to add back.
              UpperCAmelCase_     :				Optional[Any]  =     od / f'''{type_path}_results/{trainer.global_step:05d}.txt'''
              UpperCAmelCase_     :				List[str]  =     od / f'''{type_path}_generations/{trainer.global_step:05d}.txt'''
              results_file.parent.mkdir(exist_ok=_SCREAMING_SNAKE_CASE    )
              generations_file.parent.mkdir(exist_ok=_SCREAMING_SNAKE_CASE    )
         with open(_SCREAMING_SNAKE_CASE  ,'''a+'''    ) as writer:
              for key in sorted(_SCREAMING_SNAKE_CASE    ):
                   if key in ["log", "progress_bar", "preds"]:
                        continue
                   UpperCAmelCase_     :				str  =     metrics[key]
                   if isinstance(_SCREAMING_SNAKE_CASE  ,torch.Tensor    ):
                        UpperCAmelCase_     :				int  =     val.item()
                   UpperCAmelCase_     :				List[Any]  =     f'''{key}: {val:.6f}\n'''
                   writer.write(_SCREAMING_SNAKE_CASE    )

         if not save_generations:
              return

         if "preds" in metrics:
              UpperCAmelCase_     :				int  =     '''\n'''.join(metrics['''preds''']    )
              generations_file.open('''w+'''    ).write(_SCREAMING_SNAKE_CASE    )





    @rank_zero_only
    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )					->						List[Any]:
         try:
              UpperCAmelCase_     :				int  =     pl_module.model.model.num_parameters()
         except AttributeError:
              UpperCAmelCase_     :				Optional[Any]  =     pl_module.model.num_parameters()

         UpperCAmelCase_     :				Dict  =     count_trainable_parameters(_SCREAMING_SNAKE_CASE    )
         # mp stands for million parameters
         trainer.logger.log_metrics({'''n_params''': npars, '''mp''': npars / 1e6, '''grad_mp''': n_trainable_pars / 1e6}    )





    @rank_zero_only
    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )					->						int:
         save_json(pl_module.metrics  ,pl_module.metrics_save_path    )
         return self._write_logs(_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE  ,'''test'''    )





    @rank_zero_only
    def 				a__ ( self  ,_SCREAMING_SNAKE_CASE  ,_SCREAMING_SNAKE_CASE    )					->						List[str]:
         save_json(pl_module.metrics  ,pl_module.metrics_save_path    )
         # Uncommenting this will save val generations
         # return self._write_logs(trainer, pl_module, "valid")





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 lowerCamelCase__  (     _lowercase					,		_lowercase ):

     '''simple docstring'''


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







def lowerCamelCase__  (     _lowercase ):

     '''simple docstring'''


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







def lowerCamelCase__  (     _lowercase					,		_lowercase					,		_lowercase					,		_lowercase					,		_lowercase ):

     '''simple docstring'''


     if bloom_config_file == "":
          UpperCAmelCase_     :				Tuple  =     BloomConfig()
     else:
          UpperCAmelCase_     :				Optional[int]  =     BloomConfig.from_json_file(_lowercase )

     if shard_model:
          UpperCAmelCase_     :				Any  =     os.listdir(_lowercase )
          UpperCAmelCase_     :				Union[str, Any]  =     sorted(filter(lambda _lowercase : s.startswith('''layer''' ) and "model_00" in s					,		_lowercase ) )

          UpperCAmelCase_     :				Any  =     {'''weight_map''': {}, '''metadata''': {}}
          UpperCAmelCase_     :				List[str]  =     0

          UpperCAmelCase_     :				Any  =     None

          UpperCAmelCase_     :				Optional[int]  =     BloomConfig()

          for j, file in enumerate(_lowercase ):
               print('''Processing file: {}'''.format(_lowercase ) )
               UpperCAmelCase_     :				Optional[Any]  =     None

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

                    # Rename keys in the transformers names
                    UpperCAmelCase_     :				Dict  =     list(temp.keys() )
                    for key in keys:
                         UpperCAmelCase_     :				Union[str, Any]  =     temp.pop(_lowercase )

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

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

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

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

          UpperCAmelCase_     :				Tuple  =     os.listdir(_lowercase )
          UpperCAmelCase_     :				Union[str, Any]  =     sorted(filter(lambda _lowercase : s.startswith('''layer''' ) and "model_00" in s					,		_lowercase ) )

          UpperCAmelCase_     :				Any  =     None
          for i, file in enumerate(_lowercase ):
               UpperCAmelCase_     :				Optional[Any]  =     None
               for i in range(_lowercase ):
                    # load all TP files
                    UpperCAmelCase_     :				List[Any]  =     file.replace('''model_00'''					,		f'''model_0{i}''' )
                    UpperCAmelCase_     :				Optional[int]  =     torch.load(os.path.join(_lowercase					,		_lowercase )					,		map_location='''cpu''' )

                    # Rename keys in the transformers names
                    UpperCAmelCase_     :				str  =     list(temp.keys() )
                    for key in keys:
                         UpperCAmelCase_     :				Dict  =     temp.pop(_lowercase )

                    if tensors is None:
                         UpperCAmelCase_     :				int  =     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(_lowercase ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH ):
                                   tensors[key] += temp[key]
                              else:
                                   # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                                   UpperCAmelCase_     :				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
                                   UpperCAmelCase_     :				List[str]  =     torch.cat([tensors[key], temp[key]]					,		dim=_lowercase )

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

               UpperCAmelCase_     :				Tuple  =     model.load_state_dict(_lowercase					,		strict=_lowercase )
               assert not other_keys.unexpected_keys, f'''The keys {other_keys.unexpected_keys} are unexpected'''
               if missing_keys is None:
                    UpperCAmelCase_     :				Union[str, Any]  =     set(other_keys.missing_keys )
               else:
                    UpperCAmelCase_     :				Dict  =     missing_keys.intersection(set(other_keys.missing_keys ) )

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

          # Save pytorch-model
          os.makedirs(_lowercase					,		exist_ok=_lowercase )
          UpperCAmelCase_     :				str  =     pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
          UpperCAmelCase_     :				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:
               UpperCAmelCase_     :				Optional[int]  =     model.to(config.torch_dtype )
          torch.save(model.state_dict()					,		_lowercase )
          print(f'''Save configuration file to {pytorch_config_dump_path}''' )
          with open(_lowercase					,		'''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,
     )