



"""simple docstring"""






from __future__ import annotations

import random
import unittest

from transformers import TransfoXLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import (
       TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFTransfoXLForSequenceClassification,
       TFTransfoXLLMHeadModel,
       TFTransfoXLModel,
   )






class 	__snake_case	:




  def __init__(			self					:						Dict						, __lowerCAmelCase					:						Any						, ):






         """simple docstring"""






         _lowerCamelCase						:    str         =			parent
         _lowerCamelCase						:    Tuple         =			1_3
         _lowerCamelCase						:    Union[str, Any]         =			7
         _lowerCamelCase						:    Optional[Any]         =			3_0
         _lowerCamelCase						:    Dict         =			self.seq_length + self.mem_len
         _lowerCamelCase						:    List[Any]         =			1_5
         _lowerCamelCase						:    Optional[Any]         =			True
         _lowerCamelCase						:    str         =			True
         _lowerCamelCase						:    Optional[Any]         =			9_9
         _lowerCamelCase						:    List[Any]         =			[1_0, 5_0, 8_0]
         _lowerCamelCase						:    Optional[Any]         =			3_2
         _lowerCamelCase						:    Optional[Any]         =			3_2
         _lowerCamelCase						:    List[Any]         =			4
         _lowerCamelCase						:    Union[str, Any]         =			8
         _lowerCamelCase						:    List[Any]         =			1_2_8
         _lowerCamelCase						:    Union[str, Any]         =			2
         _lowerCamelCase						:    Dict         =			2
         _lowerCamelCase						:    List[str]         =			None
         _lowerCamelCase						:    List[Any]         =			1
         _lowerCamelCase						:    Union[str, Any]         =			0
         _lowerCamelCase						:    Union[str, Any]         =			3
         _lowerCamelCase						:    List[str]         =			self.vocab_size - 1
         _lowerCamelCase						:    Any         =			0.01




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size		)
         _lowerCamelCase						:    List[Any]         =			ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size		)

         _lowerCamelCase						:    Any         =			None
         if self.use_labels:
                _lowerCamelCase						:    Optional[Any]         =			ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size		)

         _lowerCamelCase						:    Optional[Any]         =			TransfoXLConfig(
             vocab_size=self.vocab_size						, mem_len=self.mem_len						, clamp_len=self.clamp_len						, cutoffs=self.cutoffs						, d_model=self.hidden_size						, d_embed=self.d_embed						, n_head=self.num_attention_heads						, d_head=self.d_head						, d_inner=self.d_inner						, div_val=self.div_val						, n_layer=self.num_hidden_layers						, eos_token_id=self.eos_token_id						, pad_token_id=self.vocab_size - 1						, init_range=self.init_range						, num_labels=self.num_labels						, )

         return (config, input_ids_a, input_ids_a, lm_labels)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         random.seed(self.seed		)
         tf.random.set_seed(self.seed		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			TFTransfoXLModel(__lowerCAmelCase		)

         _lowerCamelCase      ,   _lowerCamelCase						:    List[Any]         =			model(__lowerCAmelCase		).to_tuple()

         _lowerCamelCase						:    Optional[Any]         =			{'''input_ids''': input_ids_a, '''mems''': mems_a}

         _lowerCamelCase      ,   _lowerCamelCase						:    Optional[Any]         =			model(__lowerCAmelCase		).to_tuple()

         self.parent.assertEqual(hidden_states_a.shape						, (self.batch_size, self.seq_length, self.hidden_size)		)
         self.parent.assertEqual(hidden_states_a.shape						, (self.batch_size, self.seq_length, self.hidden_size)		)
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]						, [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						, )
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]						, [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			TFTransfoXLLMHeadModel(__lowerCAmelCase		)

         _lowerCamelCase      ,   _lowerCamelCase						:    int         =			model(__lowerCAmelCase		).to_tuple()

         _lowerCamelCase						:    str         =			{'''input_ids''': input_ids_a, '''labels''': lm_labels}
         _lowerCamelCase      ,   _lowerCamelCase						:    Dict         =			model(__lowerCAmelCase		).to_tuple()

         _lowerCamelCase      ,   _lowerCamelCase						:    Optional[Any]         =			model([input_ids_a, mems_a]		).to_tuple()

         _lowerCamelCase						:    int         =			{'''input_ids''': input_ids_a, '''mems''': mems_a, '''labels''': lm_labels}

         _lowerCamelCase      ,   _lowerCamelCase						:    Tuple         =			model(__lowerCAmelCase		).to_tuple()

         self.parent.assertEqual(lm_logits_a.shape						, (self.batch_size, self.seq_length, self.vocab_size)		)
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]						, [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						, )

         self.parent.assertEqual(lm_logits_a.shape						, (self.batch_size, self.seq_length, self.vocab_size)		)
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]						, [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			TFTransfoXLForSequenceClassification(__lowerCAmelCase		)
         _lowerCamelCase						:    List[Any]         =			model(__lowerCAmelCase		)
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			self.prepare_config_and_inputs()
         ((_lowerCamelCase)      ,   (_lowerCamelCase)      ,   (_lowerCamelCase)      ,   (_lowerCamelCase))						:    Dict         =			config_and_inputs
         _lowerCamelCase						:    Optional[Any]         =			{'''input_ids''': input_ids_a}
         return config, inputs_dict






@require_tf
class 	__snake_case	(     _lowercase  ,					_lowercase  ,					unittest.TestCase):
  snake_case__    :       Dict        =						(
      (TFTransfoXLModel, TFTransfoXLLMHeadModel, TFTransfoXLForSequenceClassification) if is_tf_available() else ()
  )
  snake_case__    :       Union[str, Any]        =						() if is_tf_available() else ()
  snake_case__    :       List[Any]        =						(
      {
          "feature-extraction": TFTransfoXLModel,
          "text-classification": TFTransfoXLForSequenceClassification,
          "text-generation": TFTransfoXLLMHeadModel,
          "zero-shot": TFTransfoXLForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  # TODO: add this test when TFTransfoXLLMHead has a linear output layer implemented
  snake_case__    :       Union[str, Any]        =						False
  snake_case__    :       Tuple        =						False
  snake_case__    :       str        =						False
  snake_case__    :       Tuple        =						False




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Optional[int]		):






         """simple docstring"""






         if pipeline_test_casse_name == "TextGenerationPipelineTests":
                # Get `ValueError: AttributeError: 'NoneType' object has no attribute 'new_ones'` or `AssertionError`.
                # `TransfoXLConfig` was never used in pipeline tests: cannot create a simple
                # tokenizer.
                return True

         return False




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			TFTransfoXLModelTester(self		)
         _lowerCamelCase						:    Union[str, Any]         =			ConfigTester(self						, config_class=__lowerCAmelCase						, d_embed=3_7		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         self.config_tester.run_common_tests()




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         self.model_tester.set_seed()
         _lowerCamelCase						:    Optional[int]         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_transfo_xl_model(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         self.model_tester.set_seed()
         _lowerCamelCase						:    Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_transfo_xl_lm_head(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_transfo_xl_for_sequence_classification(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    List[str]         =			self.model_tester.prepare_config_and_inputs_for_common()
         _lowerCamelCase						:    Optional[int]         =			[TFTransfoXLForSequenceClassification]

         for model_class in self.all_model_classes:
                _lowerCamelCase						:    Any         =			model_class(__lowerCAmelCase		)
                assert isinstance(model.get_input_embeddings()						, tf.keras.layers.Layer		)
                if model_class in list_other_models_with_output_ebd:
                       _lowerCamelCase						:    Dict         =			model.get_output_embeddings()
                       assert isinstance(__lowerCAmelCase						, tf.keras.layers.Layer		)
                       _lowerCamelCase						:    Tuple         =			model.get_bias()
                       assert name is None
                else:
                       _lowerCamelCase						:    Union[str, Any]         =			model.get_output_embeddings()
                       assert x is None
                       _lowerCamelCase						:    int         =			model.get_bias()
                       assert name is None




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         pass




  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         for model_name in TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                _lowerCamelCase						:    Dict         =			TFTransfoXLModel.from_pretrained(__lowerCAmelCase		)
                self.assertIsNotNone(__lowerCAmelCase		)




  @unittest.skip(reason='''This model doesn\'t play well with fit() due to not returning a single loss.'''		)
  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         pass






@require_tf
class 	__snake_case	(     unittest.TestCase):




  @unittest.skip('''Skip test until #12651 is resolved.'''		)
  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			TFTransfoXLLMHeadModel.from_pretrained('''transfo-xl-wt103'''		)
         # fmt: off
         _lowerCamelCase						:    Optional[int]         =			tf.convert_to_tensor([[3_3,1_2_9_7,2,1,1_0_0_9,4,1_1_0_9,1_1_7_3_9,4_7_6_2,3_5_8,5,2_5,2_4_5,2_2,1_7_0_6,1_7,2_0_0_9_8,5,3_2_1_5,2_1,3_7,1_1_1_0,3,1_3,1_0_4_1,4,2_4,6_0_3,4_9_0,2,7_1_4_7_7,2_0_0_9_8,1_0_4_4_4_7,2,2_0_9_6_1,1,2_6_0_4,4,1,3_2_9,3,6_2_2_4,8_3_1,1_6_0_0_2,2,8,6_0_3,7_8_9_6_7,2_9_5_4_6,2_3,8_0_3,2_0,2_5,4_1_6,5,8,2_3_2,4,2_7_7,6,1_8_5_5,4_6_0_1,3,2_9_5_4_6,5_4,8,3_6_0_9,5,5_7_2_1_1,4_9,4,1,2_7_7,1_8,8,1_7_5_5,1_5_6_9_1,3,3_4_1,2_5,4_1_6,6_9_3,4_2_5_7_3,7_1,1_7,4_0_1,9_4,3_1,1_7_9_1_9,2,2_9_5_4_6,7_8_7_3,1_8,1,4_3_5,2_3,1_1_0_1_1,7_5_5,5,5_1_6_7,3,7_9_8_3,9_8,8_4,2,2_9_5_4_6,3_2_6_7,8,3_6_0_9,4,1,4_8_6_5,1_0_7_5,2,6_0_8_7,7_1,6,3_4_6,8,5_8_5_4,3,2_9_5_4_6,8_2_4,1_4_0_0,1_8_6_8,2,1_9,1_6_0,2,3_1_1,8,5_4_9_6,2,2_0_9_2_0,1_7,2_5,1_5_0_9_7,3,2_4,2_4,0]]						, dtype=tf.intaa		)  # noqa: E231
         # fmt: on
         #  In 1991 , the remains of Russian Tsar Nicholas II and his family
         #  ( except for Alexei and Maria ) are discovered .
         #  The voice of Nicholas's young son , Tsarevich Alexei Nikolaevich , narrates the
         #  remainder of the story . 1883 Western Siberia ,
         #  a young Grigori Rasputin is asked by his father and a group of men to perform magic .
         #  Rasputin has a vision and denounces one of the men as a horse thief . Although his
         #  father initially slaps him for making such an accusation , Rasputin watches as the
         #  man is chased outside and beaten . Twenty years later , Rasputin sees a vision of
         #  the Virgin Mary , prompting him to become a priest . Rasputin quickly becomes famous ,
         #  with people , even a bishop , begging for his blessing . <eod> </s> <eos>

         # fmt: off
         _lowerCamelCase						:    str         =			[3_3,1_2_9_7,2,1,1_0_0_9,4,1_1_0_9,1_1_7_3_9,4_7_6_2,3_5_8,5,2_5,2_4_5,2_2,1_7_0_6,1_7,2_0_0_9_8,5,3_2_1_5,2_1,3_7,1_1_1_0,3,1_3,1_0_4_1,4,2_4,6_0_3,4_9_0,2,7_1_4_7_7,2_0_0_9_8,1_0_4_4_4_7,2,2_0_9_6_1,1,2_6_0_4,4,1,3_2_9,3,6_2_2_4,8_3_1,1_6_0_0_2,2,8,6_0_3,7_8_9_6_7,2_9_5_4_6,2_3,8_0_3,2_0,2_5,4_1_6,5,8,2_3_2,4,2_7_7,6,1_8_5_5,4_6_0_1,3,2_9_5_4_6,5_4,8,3_6_0_9,5,5_7_2_1_1,4_9,4,1,2_7_7,1_8,8,1_7_5_5,1_5_6_9_1,3,3_4_1,2_5,4_1_6,6_9_3,4_2_5_7_3,7_1,1_7,4_0_1,9_4,3_1,1_7_9_1_9,2,2_9_5_4_6,7_8_7_3,1_8,1,4_3_5,2_3,1_1_0_1_1,7_5_5,5,5_1_6_7,3,7_9_8_3,9_8,8_4,2,2_9_5_4_6,3_2_6_7,8,3_6_0_9,4,1,4_8_6_5,1_0_7_5,2,6_0_8_7,7_1,6,3_4_6,8,5_8_5_4,3,2_9_5_4_6,8_2_4,1_4_0_0,1_8_6_8,2,1_9,1_6_0,2,3_1_1,8,5_4_9_6,2,2_0_9_2_0,1_7,2_5,1_5_0_9_7,3,2_4,2_4,0,3_3,1,1_8_5_7,2,1,1_0_0_9,4,1_1_0_9,1_1_7_3_9,4_7_6_2,3_5_8,5,2_5,2_4_5,2_8,1_1_1_0,3,1_3,1_0_4_1,4,2_4,6_0_3,4_9_0,2,7_1_4_7_7,2_0_0_9_8,1_0_4_4_4_7,2,2_0_9_6_1,1,2_6_0_4,4,1,3_2_9,3,0]  # noqa: E231
         # fmt: on
         #  In 1991, the remains of Russian Tsar Nicholas II and his family (
         #  except for Alexei and Maria ) are discovered. The voice of young son,
         #  Tsarevich Alexei Nikolaevich, narrates the remainder of the story.
         #  1883 Western Siberia, a young Grigori Rasputin is asked by his father
         #  and a group of men to perform magic. Rasputin has a vision and
         #  denounces one of the men as a horse thief. Although his father initially
         #  slaps him for making such an accusation, Rasputin watches as the man
         #  is chased outside and beaten. Twenty years later, Rasputin sees a vision
         #  of the Virgin Mary, prompting him to become a priest.
         #  Rasputin quickly becomes famous, with people, even a bishop, begging for
         #  his blessing. <unk> <unk> <eos> In the 1990s, the remains of Russian Tsar
         # Nicholas II and his family were discovered. The voice of <unk> young son,
         # Tsarevich Alexei Nikolaevich, narrates the remainder of the story.<eos>

         _lowerCamelCase						:    List[Any]         =			model.generate(__lowerCAmelCase						, max_length=2_0_0						, do_sample=__lowerCAmelCase		)
         self.assertListEqual(output_ids[0].numpy().tolist()						, __lowerCAmelCase		)







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

lowerCAmelCase__      =    [
    '''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''',
]

lowerCAmelCase__      =    [
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]






def 	snake_case_  (				A_ :		int,    A_ :		Dict						):


       '''simple docstring'''





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






def 	snake_case_  (				A_ :		List[Any]						):


       '''simple docstring'''





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






def 	snake_case_  (				A_ :		str,    A_ :		Any,    A_ :		int,    A_ :		List[str],    A_ :		Any						):


       '''simple docstring'''





       if bloom_config_file == "":
              _lowerCamelCase						:    Dict         =			BloomConfig()
       else:
              _lowerCamelCase						:    Any         =			BloomConfig.from_json_file(A_						)

       if shard_model:
              _lowerCamelCase						:    Optional[int]         =			os.listdir(A_						)
              _lowerCamelCase						:    List[str]         =			sorted(filter(lambda A_						: s.startswith('''layer'''						) and "model_00" in s,    A_						)						)

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

              _lowerCamelCase						:    str         =			None

              _lowerCamelCase						:    str         =			BloomConfig()

              for j, file in enumerate(A_						):
                     print('''Processing file: {}'''.format(A_						)						)
                     _lowerCamelCase						:    List[Any]         =			None

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

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

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

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

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

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

              _lowerCamelCase						:    Optional[int]         =			os.listdir(A_						)
              _lowerCamelCase						:    Union[str, Any]         =			sorted(filter(lambda A_						: s.startswith('''layer'''						) and "model_00" in s,    A_						)						)

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

                            # Rename keys in the transformers names
                            _lowerCamelCase						:    List[Any]         =			list(temp.keys()						)
                            for key in keys:
                                   _lowerCamelCase						:    Dict         =			temp.pop(A_						)

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

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

                     _lowerCamelCase						:    List[str]         =			model.load_state_dict(A_,    strict=A_						)
                     assert not other_keys.unexpected_keys, F'''The keys {other_keys.unexpected_keys} are unexpected'''
                     if missing_keys is None:
                            _lowerCamelCase						:    Optional[Any]         =			set(other_keys.missing_keys						)
                     else:
                            _lowerCamelCase						:    Union[str, 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(A_,    exist_ok=A_						)
              _lowerCamelCase						:    Union[str, Any]         =			pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
              _lowerCamelCase						:    Tuple         =			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(),    A_						)
              print(F'''Save configuration file to {pytorch_config_dump_path}'''						)
              with open(A_,    '''w''',    encoding='''utf-8'''						) as f:
                     f.write(config.to_json_string()						)


if __name__ == "__main__":
   lowerCAmelCase__      =    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''',
   )
   lowerCAmelCase__      =    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,
   )


