
"""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 					__lowerCAmelCase       (_UpperCamelCase ,						_UpperCamelCase	):
 __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    :		List[Any]		    =   int(re.match(r'.*layer_(\d*).*' ,						_UpperCamelCase	)[1]	)
 layer_number -= 3
 return F"h.{layer_number}." + key
def 					__lowerCAmelCase       (_UpperCamelCase	):
 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    :		Dict		    =   int(bit_search.groups()[0]	)
 return bit_size // 8
def 					__lowerCAmelCase       (_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase	):
 # Construct model
 if bloom_config_file == "":
  __lowerCAmelCase    :		Optional[Any]		    =   BloomConfig()
 else:
  __lowerCAmelCase    :		str		    =   BloomConfig.from_json_file(_UpperCamelCase	)

 if shard_model:
  __lowerCAmelCase    :		Tuple		    =   os.listdir(_UpperCamelCase	)
  __lowerCAmelCase    :		str		    =   sorted(filter(lambda _UpperCamelCase	: s.startswith('layer'	) and "model_00" in s ,						_UpperCamelCase	)	)

  __lowerCAmelCase    :		Tuple		    =   {'weight_map': {}, 'metadata': {}}
  __lowerCAmelCase    :		List[Any]		    =   0

  __lowerCAmelCase    :		Optional[int]		    =   None

  __lowerCAmelCase    :		Union[str, Any]		    =   BloomConfig()

  for j, file in enumerate(_UpperCamelCase	):
   print('Processing file: {}'.format(_UpperCamelCase	)	)
   __lowerCAmelCase    :		Tuple		    =   None

   for i in range(_UpperCamelCase	):
    # load all TP files
    __lowerCAmelCase    :		int		    =   file.replace('model_00' ,						F"model_0{i}"	)
    __lowerCAmelCase    :		Optional[Any]		    =   torch.load(os.path.join(_UpperCamelCase ,						_UpperCamelCase	) ,						map_location='cpu'	)

    # Rename keys in the transformers names
    __lowerCAmelCase    :		Union[str, Any]		    =   list(temp.keys()	)
    for key in keys:
     __lowerCAmelCase    :		Optional[int]		    =   temp.pop(_UpperCamelCase	)

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

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

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

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

    # Rename keys in the transformers names
    __lowerCAmelCase    :		Optional[int]		    =   list(temp.keys()	)
    for key in keys:
     __lowerCAmelCase    :		int		    =   temp.pop(_UpperCamelCase	)

    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(_UpperCamelCase	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
       tensors[key] += temp[key]
      else:
       # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
       __lowerCAmelCase    :		Union[str, 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    :		str		    =   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    :		List[str]		    =   tensors[key] / pretraining_tp

   __lowerCAmelCase    :		int		    =   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    :		Dict		    =   set(other_keys.missing_keys	)
   else:
    __lowerCAmelCase    :		List[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    :		int		    =   pytorch_dump_folder_path + '/' + WEIGHTS_NAME
  __lowerCAmelCase    :		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__":
       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,
       )

"""simple docstring"""







import inspect
import os
import re

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
lowerCamelCase__           =	"""src/transformers"""


# This is to make sure the transformers module imported is the one in the repo.
lowerCamelCase__           =	direct_transformers_import(PATH_TO_TRANSFORMERS)

lowerCamelCase__           =	transformers.models.auto.configuration_auto.CONFIG_MAPPING

lowerCamelCase__           =	{
    # used to compute the property `self.chunk_length`
    """EncodecConfig""": ["""overlap"""],
    # used as `self.bert_model = BertModel(config, ...)`
    """DPRConfig""": True,
    # not used in modeling files, but it's an important information
    """FSMTConfig""": ["""langs"""],
    # used internally in the configuration class file
    """GPTNeoConfig""": ["""attention_types"""],
    # used internally in the configuration class file
    """EsmConfig""": ["""is_folding_model"""],
    # used during training (despite we don't have training script for these models yet)
    """Mask2FormerConfig""": ["""ignore_value"""],
    # `ignore_value` used during training (despite we don't have training script for these models yet)
    # `norm` used in conversion script (despite not using in the modeling file)
    """OneFormerConfig""": ["""ignore_value""", """norm"""],
    # used during preprocessing and collation, see `collating_graphormer.py`
    """GraphormerConfig""": ["""spatial_pos_max"""],
    # used internally in the configuration class file
    """T5Config""": ["""feed_forward_proj"""],
    # used internally in the configuration class file
    # `tokenizer_class` get default value `T5Tokenizer` intentionally
    """MT5Config""": ["""feed_forward_proj""", """tokenizer_class"""],
    """UMT5Config""": ["""feed_forward_proj""", """tokenizer_class"""],
    # used internally in the configuration class file
    """LongT5Config""": ["""feed_forward_proj"""],
    # used internally in the configuration class file
    """SwitchTransformersConfig""": ["""feed_forward_proj"""],
    # having default values other than `1e-5` - we can't fix them without breaking
    """BioGptConfig""": ["""layer_norm_eps"""],
    # having default values other than `1e-5` - we can't fix them without breaking
    """GLPNConfig""": ["""layer_norm_eps"""],
    # having default values other than `1e-5` - we can't fix them without breaking
    """SegformerConfig""": ["""layer_norm_eps"""],
    # having default values other than `1e-5` - we can't fix them without breaking
    """CvtConfig""": ["""layer_norm_eps"""],
    # having default values other than `1e-5` - we can't fix them without breaking
    """PerceiverConfig""": ["""layer_norm_eps"""],
    # used internally to calculate the feature size
    """InformerConfig""": ["""num_static_real_features""", """num_time_features"""],
    # used internally to calculate the feature size
    """TimeSeriesTransformerConfig""": ["""num_static_real_features""", """num_time_features"""],
    # used internally to calculate the feature size
    """AutoformerConfig""": ["""num_static_real_features""", """num_time_features"""],
    # used internally to calculate `mlp_dim`
    """SamVisionConfig""": ["""mlp_ratio"""],
    # For (head) training, but so far not implemented
    """ClapAudioConfig""": ["""num_classes"""],
    # Not used, but providing useful information to users
    """SpeechT5HifiGanConfig""": ["""sampling_rate"""],
}


# TODO (ydshieh): Check the failing cases, try to fix them or move some cases to the above block once we are sure
SPECIAL_CASES_TO_ALLOW.update(
    {
        """CLIPSegConfig""": True,
        """DeformableDetrConfig""": True,
        """DetaConfig""": True,
        """DinatConfig""": True,
        """DonutSwinConfig""": True,
        """EfficientFormerConfig""": True,
        """FSMTConfig""": True,
        """JukeboxConfig""": True,
        """LayoutLMv2Config""": True,
        """MaskFormerSwinConfig""": True,
        """MT5Config""": True,
        """NatConfig""": True,
        """OneFormerConfig""": True,
        """PerceiverConfig""": True,
        """RagConfig""": True,
        """SpeechT5Config""": True,
        """SwinConfig""": True,
        """Swin2SRConfig""": True,
        """Swinv2Config""": True,
        """SwitchTransformersConfig""": True,
        """TableTransformerConfig""": True,
        """TapasConfig""": True,
        """TransfoXLConfig""": True,
        """UniSpeechConfig""": True,
        """UniSpeechSatConfig""": True,
        """WavLMConfig""": True,
        """WhisperConfig""": True,
        # TODO: @Arthur (for `alignment_head` and `alignment_layer`)
        """JukeboxPriorConfig""": True,
        # TODO: @Younes (for `is_decoder`)
        """Pix2StructTextConfig""": True,
    }
)
def 					__lowerCAmelCase       (_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase	):
	__lowerCAmelCase    :		Optional[int]		    =   False
	for attribute in attributes:
		for modeling_source in source_strings:
			# check if we can find `config.xxx`, `getattr(config, "xxx", ...)` or `getattr(self.config, "xxx", ...)`
			if (
			    F"config.{attribute}" in modeling_source
			    or F"getattr(config, \"{attribute}\"" in modeling_source
			    or F"getattr(self.config, \"{attribute}\"" in modeling_source
			):
				__lowerCAmelCase    :		Union[str, Any]		    =   True
			# Deal with multi-line cases
			elif (
			    re.search(
			        rF"getattr[ \t\v\n\r\f]*\([ \t\v\n\r\f]*(self\.)?config,[ \t\v\n\r\f]*\"{attribute}\"" ,						_UpperCamelCase ,						)
			    is not None
			):
				__lowerCAmelCase    :		Optional[Any]		    =   True
			# `SequenceSummary` is called with `SequenceSummary(config)`
			elif attribute in [
			    "summary_type",
			    "summary_use_proj",
			    "summary_activation",
			    "summary_last_dropout",
			    "summary_proj_to_labels",
			    "summary_first_dropout",
			]:
				if "SequenceSummary" in modeling_source:
					__lowerCAmelCase    :		int		    =   True
			if attribute_used:
				break
		if attribute_used:
			break

    # common and important attributes, even if they do not always appear in the modeling files
	__lowerCAmelCase    :		Optional[Any]		    =   [
	    'bos_index',
	    'eos_index',
	    'pad_index',
	    'unk_index',
	    'mask_index',
	    'image_size',
	    'use_cache',
	    'out_features',
	    'out_indices',
	]
	__lowerCAmelCase    :		Tuple		    =   ['encoder_no_repeat_ngram_size']

	# Special cases to be allowed
	__lowerCAmelCase    :		Dict		    =   True
	if not attribute_used:
		__lowerCAmelCase    :		int		    =   False
		for attribute in attributes:
			# Allow if the default value in the configuration class is different from the one in `PretrainedConfig`
			if attribute in ["is_encoder_decoder"] and default_value is True:
				__lowerCAmelCase    :		int		    =   True
			elif attribute in ["tie_word_embeddings"] and default_value is False:
				__lowerCAmelCase    :		Dict		    =   True

			# Allow cases without checking the default value in the configuration class
			elif attribute in attributes_to_allow + attributes_used_in_generation:
				__lowerCAmelCase    :		Tuple		    =   True
			elif attribute.endswith('_token_id'	):
				__lowerCAmelCase    :		int		    =   True

			# configuration class specific cases
			if not case_allowed:
				__lowerCAmelCase    :		Dict		    =   SPECIAL_CASES_TO_ALLOW.get(config_class.__name__ ,						[]	)
				__lowerCAmelCase    :		Any		    =   allowed_cases is True or attribute in allowed_cases

	return attribute_used or case_allowed
def 					__lowerCAmelCase       (_UpperCamelCase	):
	__lowerCAmelCase    :		List[str]		    =   dict(inspect.signature(config_class.__init__	).parameters	)
	__lowerCAmelCase    :		Tuple		    =   [x for x in list(signature.keys()	) if x not in ['self', 'kwargs']]
	__lowerCAmelCase    :		Tuple		    =   [signature[param].default for param in parameter_names]

	# If `attribute_map` exists, an attribute can have different names to be used in the modeling files, and as long
	# as one variant is used, the test should pass
	__lowerCAmelCase    :		List[Any]		    =   {}
	if len(config_class.attribute_map	) > 0:
		__lowerCAmelCase    :		Union[str, Any]		    =   {v: k for k, v in config_class.attribute_map.items()}

	# Get the path to modeling source files
	__lowerCAmelCase    :		Dict		    =   inspect.getsourcefile(_UpperCamelCase	)
	__lowerCAmelCase    :		List[str]		    =   os.path.dirname(_UpperCamelCase	)
	# Let's check against all frameworks: as long as one framework uses an attribute, we are good.
	__lowerCAmelCase    :		Tuple		    =   [os.path.join(_UpperCamelCase ,						_UpperCamelCase	) for fn in os.listdir(_UpperCamelCase	) if fn.startswith('modeling_'	)]

	# Get the source code strings
	__lowerCAmelCase    :		str		    =   []
	for path in modeling_paths:
		if os.path.isfile(_UpperCamelCase	):
			with open(_UpperCamelCase	) as fp:
				modeling_sources.append(fp.read()	)

	__lowerCAmelCase    :		Optional[int]		    =   []
	for config_param, default_value in zip(_UpperCamelCase ,						_UpperCamelCase	):
		# `attributes` here is all the variant names for `config_param`
		__lowerCAmelCase    :		Optional[int]		    =   [config_param]
		# some configuration classes have non-empty `attribute_map`, and both names could be used in the
		# corresponding modeling files. As long as one of them appears, it is fine.
		if config_param in reversed_attribute_map:
			attributes.append(reversed_attribute_map[config_param]	)

		if not check_attribute_being_used(_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase	):
			unused_attributes.append(attributes[0]	)

	return sorted(_UpperCamelCase	)
def 					__lowerCAmelCase       ():
	__lowerCAmelCase    :		List[str]		    =   {}
	for _config_class in list(CONFIG_MAPPING.values()	):
		# Skip deprecated models
		if "models.deprecated" in _config_class.__module__:
			continue
		# Some config classes are not in `CONFIG_MAPPING` (e.g. `CLIPVisionConfig`, `Blip2VisionConfig`, etc.)
		__lowerCAmelCase    :		int		    =   [
		    cls
		    for name, cls in inspect.getmembers(
		        inspect.getmodule(_config_class	) ,						lambda _UpperCamelCase	: inspect.isclass(_UpperCamelCase	)
		        and issubclass(_UpperCamelCase ,						_UpperCamelCase	)
		        and inspect.getmodule(_UpperCamelCase	) == inspect.getmodule(_config_class	) ,						)
		]
		for config_class in config_classes_in_module:
			__lowerCAmelCase    :		List[Any]		    =   check_config_attributes_being_used(_UpperCamelCase	)
			if len(_UpperCamelCase	) > 0:
				__lowerCAmelCase    :		str		    =   unused_attributes

	if len(_UpperCamelCase	) > 0:
		__lowerCAmelCase    :		int		    =   'The following configuration classes contain unused attributes in the corresponding modeling files:\n'
		for name, attributes in configs_with_unused_attributes.items():
			error += F"{name}: {attributes}\n"

		raise ValueError(_UpperCamelCase	)


if __name__ == "__main__":
							check_config_attributes()