
'''simple docstring'''


import torch
from torch import nn


class 						lowercase__	(  nn.Module   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=1 ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      super().__init__()

      UpperCamelCase       =  n_token
      UpperCamelCase       =  d_embed
      UpperCamelCase       =  d_proj

      UpperCamelCase       =  cutoffs + [n_token]
      UpperCamelCase       =  [0] + self.cutoffs
      UpperCamelCase       =  div_val

      UpperCamelCase       =  self.cutoffs[0]
      UpperCamelCase       =  len(self.cutoffs   ) - 1
      UpperCamelCase       =  self.shortlist_size + self.n_clusters

      if self.n_clusters > 0:
         UpperCamelCase       =  nn.Parameter(torch.zeros(self.n_clusters ,   self.d_embed   )   )
         UpperCamelCase       =  nn.Parameter(torch.zeros(self.n_clusters   )   )

      UpperCamelCase       =  nn.ModuleList()
      UpperCamelCase       =  nn.ParameterList()

      if div_val == 1:
         for i in range(len(self.cutoffs   )   ):
            if d_proj != d_embed:
               self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCamelCase__ ,   lowerCamelCase__   )   )   )
            else:
               self.out_projs.append(lowerCamelCase__   )

         self.out_layers.append(nn.Linear(lowerCamelCase__ ,   lowerCamelCase__   )   )
      else:
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            UpperCamelCase       =  d_embed // (div_val**i)

            self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCamelCase__ ,   lowerCamelCase__   )   )   )

            self.out_layers.append(nn.Linear(lowerCamelCase__ ,   r_idx - l_idx   )   )

      UpperCamelCase       =  keep_order







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      if proj is None:
         UpperCamelCase       =  nn.functional.linear(lowerCamelCase__ ,   lowerCamelCase__ ,   bias=lowerCamelCase__   )
      else:
         # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
         UpperCamelCase       =  nn.functional.linear(lowerCamelCase__ ,   proj.t().contiguous()   )
         UpperCamelCase       =  nn.functional.linear(lowerCamelCase__ ,   lowerCamelCase__ ,   bias=lowerCamelCase__   )
         # else:
         #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
         #     if bias is not None:
         #         logit = logit + bias

      return logit







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      if labels is not None:
         # Shift so that tokens < n predict n
         UpperCamelCase       =  hidden[..., :-1, :].contiguous()
         UpperCamelCase       =  labels[..., 1:].contiguous()
         UpperCamelCase       =  hidden.view(-1 ,   hidden.size(-1   )   )
         UpperCamelCase       =  labels.view(-1   )
         if hidden.size(0   ) != labels.size(0   ):
            raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''   )
      else:
         UpperCamelCase       =  hidden.view(-1 ,   hidden.size(-1   )   )

      if self.n_clusters == 0:
         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   self.out_layers[0].weight ,   self.out_layers[0].bias ,   self.out_projs[0]   )
         if labels is not None:
            UpperCamelCase       =  labels != -1_0_0
            UpperCamelCase       =  torch.zeros_like(lowerCamelCase__ ,   dtype=hidden.dtype ,   device=hidden.device   )
            UpperCamelCase       =  (
                -nn.functional.log_softmax(lowerCamelCase__ ,   dim=-1   )[mask].gather(1 ,   labels[mask].unsqueeze(1   )   ).squeeze(1   )
            )
         else:
            UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=-1   )
      else:
         # construct weights and biases
         UpperCamelCase     , UpperCamelCase       =  [], []
         for i in range(len(self.cutoffs   )   ):
            if self.div_val == 1:
               UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
               UpperCamelCase       =  self.out_layers[0].weight[l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0].bias[l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i].weight
               UpperCamelCase       =  self.out_layers[i].bias

            if i == 0:
               UpperCamelCase       =  torch.cat([weight_i, self.cluster_weight] ,   dim=0   )
               UpperCamelCase       =  torch.cat([bias_i, self.cluster_bias] ,   dim=0   )

            weights.append(lowerCamelCase__   )
            biases.append(lowerCamelCase__   )

         UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[0], biases[0], self.out_projs[0]

         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )
         UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )

         if labels is None:
            UpperCamelCase       =  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
         else:
            UpperCamelCase       =  torch.zeros_like(lowerCamelCase__ ,   dtype=hidden.dtype ,   device=hidden.device   )

         UpperCamelCase       =  0
         UpperCamelCase       =  [0] + self.cutoffs
         for i in range(len(lowerCamelCase__   ) - 1   ):
            UpperCamelCase     , UpperCamelCase       =  cutoff_values[i], cutoff_values[i + 1]

            if labels is not None:
               UpperCamelCase       =  (labels >= l_idx) & (labels < r_idx)
               UpperCamelCase       =  mask_i.nonzero().squeeze()

               if indices_i.numel() == 0:
                  continue

               UpperCamelCase       =  labels.index_select(0 ,   lowerCamelCase__   ) - l_idx
               UpperCamelCase       =  head_logprob.index_select(0 ,   lowerCamelCase__   )
               UpperCamelCase       =  hidden.index_select(0 ,   lowerCamelCase__   )
            else:
               UpperCamelCase       =  hidden

            if i == 0:
               if labels is not None:
                  UpperCamelCase       =  head_logprob_i.gather(1 ,   target_i[:, None]   ).squeeze(1   )
               else:
                  UpperCamelCase       =  head_logprob[:, : self.cutoffs[0]]
            else:
               UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[i], biases[i], self.out_projs[i]

               UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )
               UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )
               UpperCamelCase       =  self.cutoffs[0] + i - 1  # No probability for the head cluster
               if labels is not None:
                  UpperCamelCase       =  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                      1 ,   target_i[:, None]   ).squeeze(1   )
               else:
                  UpperCamelCase       =  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                  UpperCamelCase       =  logprob_i

            if labels is not None:
               if (hasattr(self ,   '''keep_order'''   ) and self.keep_order) or keep_order:
                  out.index_copy_(0 ,   lowerCamelCase__ ,   -logprob_i   )
               else:
                  out[offset : offset + logprob_i.size(0   )].copy_(-logprob_i   )
               offset += logprob_i.size(0   )

      return out







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if self.n_clusters == 0:
         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   self.out_layers[0].weight ,   self.out_layers[0].bias ,   self.out_projs[0]   )
         return nn.functional.log_softmax(lowerCamelCase__ ,   dim=-1   )
      else:
         # construct weights and biases
         UpperCamelCase     , UpperCamelCase       =  [], []
         for i in range(len(self.cutoffs   )   ):
            if self.div_val == 1:
               UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
               UpperCamelCase       =  self.out_layers[0].weight[l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0].bias[l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i].weight
               UpperCamelCase       =  self.out_layers[i].bias

            if i == 0:
               UpperCamelCase       =  torch.cat([weight_i, self.cluster_weight] ,   dim=0   )
               UpperCamelCase       =  torch.cat([bias_i, self.cluster_bias] ,   dim=0   )

            weights.append(lowerCamelCase__   )
            biases.append(lowerCamelCase__   )

         UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[0], biases[0], self.out_projs[0]
         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )

         UpperCamelCase       =  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
         UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )

         UpperCamelCase       =  [0] + self.cutoffs
         for i in range(len(lowerCamelCase__   ) - 1   ):
            UpperCamelCase     , UpperCamelCase       =  cutoff_values[i], cutoff_values[i + 1]

            if i == 0:
               UpperCamelCase       =  head_logprob[:, : self.cutoffs[0]]
            else:
               UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[i], biases[i], self.out_projs[i]

               UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )
               UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )

               UpperCamelCase       =  head_logprob[:, -i] + tail_logprob_i
               UpperCamelCase       =  logprob_i

         return out







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

snake_case_ :  Tuple     =					[
    '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',
]

snake_case_ :  Union[str, Any]     =					[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]



def  __snake_case     (    _UpperCAmelCase					:       List[str],      _UpperCAmelCase					:       List[str]):
   UpperCamelCase       =  {
       '''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       =  int(re.match(R'''.*layer_(\d*).*''',      _UpperCAmelCase)[1])
   layer_number -= 3
   return f'h.{layer_number}.' + key



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



def  __snake_case     (    _UpperCAmelCase					:       int,      _UpperCAmelCase					:       Optional[Any],      _UpperCAmelCase					:       str,      _UpperCAmelCase					:       List[Any],      _UpperCAmelCase					:       Optional[int]):
   # Construct model
   if bloom_config_file == "":
      UpperCamelCase       =  BloomConfig()
   else:
      UpperCamelCase       =  BloomConfig.from_json_file(_UpperCAmelCase)

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

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

      UpperCamelCase       =  None

      UpperCamelCase       =  BloomConfig()

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

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

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

            if tensors is None:
               UpperCamelCase       =  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
                     UpperCamelCase       =  1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN) else 0
                     # We concatenate these weights accross TP ranks
                     UpperCamelCase       =  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):
               UpperCamelCase       =  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():
            UpperCamelCase       =  tensors[key]
            total_size += value.numel() * get_dtype_size(value.dtype)
            if key not in index_dict["weight_map"]:
               UpperCamelCase       =  '''pytorch_model_{}-of-{}.bin'''.format(
                   str(j + 1).zfill(5),      str(len(_UpperCAmelCase)).zfill(5))

      UpperCamelCase       =  BloomConfig()
      UpperCamelCase       =  pytorch_dump_folder_path + '''/''' + CONFIG_NAME
      UpperCamelCase       =  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:
         UpperCamelCase       =  json.dumps(_UpperCAmelCase,      indent=2,      sort_keys=_UpperCAmelCase) + '''\n'''
         f.write(_UpperCAmelCase)
   else:
      UpperCamelCase       =  BloomModel(_UpperCAmelCase)

      UpperCamelCase       =  os.listdir(_UpperCAmelCase)
      UpperCamelCase       =  sorted(filter(lambda _UpperCAmelCase: s.startswith('''layer''') and "model_00" in s,      _UpperCAmelCase))

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

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

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

         UpperCamelCase       =  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:
            UpperCamelCase       =  set(other_keys.missing_keys)
         else:
            UpperCamelCase       =  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)
      UpperCamelCase       =  pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
      UpperCamelCase       =  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       =  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__":
   snake_case_ :  Tuple     =					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',
   )
   snake_case_ :  List[str]     =					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,
   )





