
import unittest

from transformers import PegasusTokenizer, PegasusTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, require_torch, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


snake_case_ :  Tuple     =					get_tests_dir('fixtures/test_sentencepiece_no_bos.model')


@require_sentencepiece
@require_tokenizers
class 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   _snake_case		   =							PegasusTokenizer
   _snake_case		   =							PegasusTokenizerFast
   _snake_case		   =							True
   _snake_case		   =							True







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      super().setUp()

      # We have a SentencePiece fixture for testing
      UpperCamelCase       =  PegasusTokenizer(lowerCamelCase__   )
      tokenizer.save_pretrained(self.tmpdirname   )







   @cached_property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return PegasusTokenizer.from_pretrained('''google/pegasus-large'''   )







   def       UpperCAmelCase				(			self ,   **lowerCamelCase__   ):

      '''simple docstring'''







      return PegasusTokenizer.from_pretrained(self.tmpdirname ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return ("This is a test", "This is a test")







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  '''</s>'''
      UpperCamelCase       =  1

      self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowerCamelCase__   ) ,   lowerCamelCase__   )
      self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowerCamelCase__   ) ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  list(self.get_tokenizer().get_vocab().keys()   )

      self.assertEqual(vocab_keys[0] ,   '''<pad>'''   )
      self.assertEqual(vocab_keys[1] ,   '''</s>'''   )
      self.assertEqual(vocab_keys[-1] ,   '''v'''   )
      self.assertEqual(len(lowerCamelCase__   ) ,   1_1_0_3   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      self.assertEqual(self.get_tokenizer().vocab_size ,   1_1_0_3   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.rust_tokenizer_class.from_pretrained(self.tmpdirname   )
      UpperCamelCase       =  self.tokenizer_class.from_pretrained(self.tmpdirname   )
      UpperCamelCase       =  (
          '''Let\'s see which <unk> is the better <unk_token_11> one <mask_1> It seems like this <mask_2> was important'''
          ''' </s> <pad> <pad> <pad>'''
      )
      UpperCamelCase       =  rust_tokenizer([raw_input_str] ,   return_tensors=lowerCamelCase__ ,   add_special_tokens=lowerCamelCase__   ).input_ids[0]
      UpperCamelCase       =  py_tokenizer([raw_input_str] ,   return_tensors=lowerCamelCase__ ,   add_special_tokens=lowerCamelCase__   ).input_ids[0]
      self.assertListEqual(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self._large_tokenizer
      # <mask_1> masks whole sentence while <mask_2> masks single word
      UpperCamelCase       =  '''<mask_1> To ensure a <mask_2> flow of bank resolutions.'''
      UpperCamelCase       =  [2, 4_1_3, 6_1_5, 1_1_4, 3, 1_9_7_1, 1_1_3, 1_6_7_9, 1_0_7_1_0, 1_0_7, 1]
      UpperCamelCase       =  tokenizer([raw_input_str] ,   return_tensors=lowerCamelCase__   ).input_ids[0]
      self.assertListEqual(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self._large_tokenizer
      # The tracebacks for the following asserts are **better** without messages or self.assertEqual
      assert tokenizer.vocab_size == 9_6_1_0_3
      assert tokenizer.pad_token_id == 0
      assert tokenizer.eos_token_id == 1
      assert tokenizer.offset == 1_0_3
      assert tokenizer.unk_token_id == tokenizer.offset + 2 == 1_0_5
      assert tokenizer.unk_token == "<unk>"
      assert tokenizer.model_max_length == 1_0_2_4
      UpperCamelCase       =  '''To ensure a smooth flow of bank resolutions.'''
      UpperCamelCase       =  [4_1_3, 6_1_5, 1_1_4, 2_2_9_1, 1_9_7_1, 1_1_3, 1_6_7_9, 1_0_7_1_0, 1_0_7, 1]
      UpperCamelCase       =  tokenizer([raw_input_str] ,   return_tensors=lowerCamelCase__   ).input_ids[0]
      self.assertListEqual(lowerCamelCase__ ,   lowerCamelCase__   )
      assert tokenizer.convert_ids_to_tokens([0, 1, 2, 3]   ) == ["<pad>", "</s>", "<mask_1>", "<mask_2>"]







   @require_torch
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  ['''This is going to be way too long.''' * 1_5_0, '''short example''']
      UpperCamelCase       =  ['''not super long but more than 5 tokens''', '''tiny''']
      UpperCamelCase       =  self._large_tokenizer(lowerCamelCase__ ,   padding=lowerCamelCase__ ,   truncation=lowerCamelCase__ ,   return_tensors='''pt'''   )
      UpperCamelCase       =  self._large_tokenizer(
          text_target=lowerCamelCase__ ,   max_length=5 ,   padding=lowerCamelCase__ ,   truncation=lowerCamelCase__ ,   return_tensors='''pt'''   )

      assert batch.input_ids.shape == (2, 1_0_2_4)
      assert batch.attention_mask.shape == (2, 1_0_2_4)
      assert targets["input_ids"].shape == (2, 5)
      assert len(lowerCamelCase__   ) == 2  # input_ids, attention_mask.







   @slow
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  {'''input_ids''': [[3_8_9_7_9, 1_4_3, 1_8_4_8_5, 6_0_6, 1_3_0, 2_6_6_6_9, 8_7_6_8_6, 1_2_1, 5_4_1_8_9, 1_1_2_9, 1_1_1, 2_6_6_6_9, 8_7_6_8_6, 1_2_1, 9_1_1_4, 1_4_7_8_7, 1_2_1, 1_3_2_4_9, 1_5_8, 5_9_2, 9_5_6, 1_2_1, 1_4_6_2_1, 3_1_5_7_6, 1_4_3, 6_2_6_1_3, 1_0_8, 9_6_8_8, 9_3_0, 4_3_4_3_0, 1_1_5_6_2, 6_2_6_1_3, 3_0_4, 1_0_8, 1_1_4_4_3, 8_9_7, 1_0_8, 9_3_1_4, 1_7_4_1_5, 6_3_3_9_9, 1_0_8, 1_1_4_4_3, 7_6_1_4, 1_8_3_1_6, 1_1_8, 4_2_8_4, 7_1_4_8, 1_2_4_3_0, 1_4_3, 1_4_0_0, 2_5_7_0_3, 1_5_8, 1_1_1, 4_2_8_4, 7_1_4_8, 1_1_7_7_2, 1_4_3, 2_1_2_9_7, 1_0_6_4, 1_5_8, 1_2_2, 2_0_4, 3_5_0_6, 1_7_5_4, 1_1_3_3, 1_4_7_8_7, 1_5_8_1, 1_1_5, 3_3_2_2_4, 4_4_8_2, 1_1_1, 1_3_5_5, 1_1_0, 2_9_1_7_3, 3_1_7, 5_0_8_3_3, 1_0_8, 2_0_1_4_7, 9_4_6_6_5, 1_1_1, 7_7_1_9_8, 1_0_7, 1], [1_1_0, 6_2_6_1_3, 1_1_7, 6_3_8, 1_1_2, 1_1_3_3, 1_2_1, 2_0_0_9_8, 1_3_5_5, 7_9_0_5_0, 1_3_8_7_2, 1_3_5, 1_5_9_6, 5_3_5_4_1, 1_3_5_2, 1_4_1, 1_3_0_3_9, 5_5_4_2, 1_2_4, 3_0_2, 5_1_8, 1_1_1, 2_6_8, 2_9_5_6, 1_1_5, 1_4_9, 4_4_2_7, 1_0_7, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1_3_9, 1_2_3_5, 2_7_9_9, 1_8_2_8_9, 1_7_7_8_0, 2_0_4, 1_0_9, 9_4_7_4, 1_2_9_6, 1_0_7, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''attention_mask''': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
      # fmt: on

      self.tokenizer_integration_test_util(
          expected_encoding=lowerCamelCase__ ,   model_name='''google/bigbird-pegasus-large-arxiv''' ,   revision='''ba85d0851d708441f91440d509690f1ab6353415''' ,   )




@require_sentencepiece
@require_tokenizers
class 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   _snake_case		   =							PegasusTokenizer
   _snake_case		   =							PegasusTokenizerFast
   _snake_case		   =							True
   _snake_case		   =							True







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      super().setUp()

      # We have a SentencePiece fixture for testing
      UpperCamelCase       =  PegasusTokenizer(lowerCamelCase__ ,   offset=0 ,   mask_token_sent=lowerCamelCase__ ,   mask_token='''[MASK]'''   )
      tokenizer.save_pretrained(self.tmpdirname   )







   @cached_property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return PegasusTokenizer.from_pretrained('''google/bigbird-pegasus-large-arxiv'''   )







   def       UpperCAmelCase				(			self ,   **lowerCamelCase__   ):

      '''simple docstring'''







      return PegasusTokenizer.from_pretrained(self.tmpdirname ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      return ("This is a test", "This is a test")







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.rust_tokenizer_class.from_pretrained(self.tmpdirname   )
      UpperCamelCase       =  self.tokenizer_class.from_pretrained(self.tmpdirname   )
      UpperCamelCase       =  (
          '''Let\'s see which <unk> is the better <unk_token> one [MASK] It seems like this [MASK] was important </s>'''
          ''' <pad> <pad> <pad>'''
      )
      UpperCamelCase       =  rust_tokenizer([raw_input_str] ,   return_tensors=lowerCamelCase__ ,   add_special_tokens=lowerCamelCase__   ).input_ids[0]
      UpperCamelCase       =  py_tokenizer([raw_input_str] ,   return_tensors=lowerCamelCase__ ,   add_special_tokens=lowerCamelCase__   ).input_ids[0]
      self.assertListEqual(lowerCamelCase__ ,   lowerCamelCase__   )







   @require_torch
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  ['''This is going to be way too long.''' * 1_0_0_0, '''short example''']
      UpperCamelCase       =  ['''not super long but more than 5 tokens''', '''tiny''']
      UpperCamelCase       =  self._large_tokenizer(lowerCamelCase__ ,   padding=lowerCamelCase__ ,   truncation=lowerCamelCase__ ,   return_tensors='''pt'''   )
      UpperCamelCase       =  self._large_tokenizer(
          text_target=lowerCamelCase__ ,   max_length=5 ,   padding=lowerCamelCase__ ,   truncation=lowerCamelCase__ ,   return_tensors='''pt'''   )

      assert batch.input_ids.shape == (2, 4_0_9_6)
      assert batch.attention_mask.shape == (2, 4_0_9_6)
      assert targets["input_ids"].shape == (2, 5)
      assert len(lowerCamelCase__   ) == 2  # input_ids, attention_mask.







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  (
          '''This is an example string that is used to test the original TF implementation against the HF'''
          ''' implementation'''
      )

      UpperCamelCase       =  self._large_tokenizer(lowerCamelCase__   ).input_ids

      self.assertListEqual(
          lowerCamelCase__ ,   [1_8_2, 1_1_7, 1_4_2, 5_8_7, 4_2_1_1, 1_2_0, 1_1_7, 2_6_3, 1_1_2, 8_0_4, 1_0_9, 8_5_6, 2_5_0_1_6, 3_1_3_7, 4_6_4, 1_0_9, 2_6_9_5_5, 3_1_3_7, 1] ,   )







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





