



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	:   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",
]

_lowerCAmelCase	:   Optional[int]    =   [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__					)	->     int:



				"""simple docstring"""





				lowerCAmelCase__		  =  {
				    '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__		  =  int(re.match(R'.*layer_(\d*).*'    ,			snake_case__					)[1]					)
				layer_number -= 3
				return f'h.{layer_number}.' + key







def     UpperCAmelCase_		(					snake_case__					)	->     Tuple:



				"""simple docstring"""





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







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__    ,			snake_case__					)	->     List[str]:



				"""simple docstring"""





				if bloom_config_file == "":
								lowerCAmelCase__		  =  BloomConfig()
				else:
								lowerCAmelCase__		  =  BloomConfig.from_json_file(snake_case__					)

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

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

								lowerCAmelCase__		  =  None

								lowerCAmelCase__		  =  BloomConfig()

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

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

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

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

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

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

								lowerCAmelCase__		  =  BloomConfig()
								lowerCAmelCase__		  =  pytorch_dump_folder_path + '/' + CONFIG_NAME
								lowerCAmelCase__		  =  total_size
								with open(snake_case__    ,			'w'    ,			encoding='utf-8'					) as f:
												f.write(config.to_json_string()					)
								with open(os.path.join(snake_case__    ,			WEIGHTS_NAME + '.index.json'					)    ,			'w'    ,			encoding='utf-8'					) as f:
												lowerCAmelCase__		  =  json.dumps(snake_case__    ,			indent=2    ,			sort_keys=snake_case__					) + '\n'
												f.write(snake_case__					)
				else:
								lowerCAmelCase__		  =  BloomModel(snake_case__					)

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

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

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

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

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

												lowerCAmelCase__		  =  model.load_state_dict(snake_case__    ,			strict=snake_case__					)
												assert not other_keys.unexpected_keys, f'The keys {other_keys.unexpected_keys} are unexpected'
												if missing_keys is None:
																lowerCAmelCase__		  =  set(other_keys.missing_keys					)
												else:
																lowerCAmelCase__		  =  missing_keys.intersection(set(other_keys.missing_keys					)					)

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

								# Save pytorch-model
								os.makedirs(snake_case__    ,			exist_ok=snake_case__					)
								lowerCAmelCase__		  =  pytorch_dump_folder_path + '/' + WEIGHTS_NAME
								lowerCAmelCase__		  =  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__		  =  model.to(config.torch_dtype					)
								torch.save(model.state_dict()    ,			snake_case__					)
								print(f'Save configuration file to {pytorch_config_dump_path}'					)
								with open(snake_case__    ,			'w'    ,			encoding='utf-8'					) as f:
												f.write(config.to_json_string()					)


if __name__ == "__main__":
	_lowerCAmelCase	:   int    =   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	:   Optional[int]    =   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,
	)





import itertools
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Union

import pandas as pd
import pyarrow as pa

import datasets
import datasets.config
from datasets.features.features import require_storage_cast
from datasets.table import table_cast
from datasets.utils.py_utils import Literal


_lowerCAmelCase	:   Any    =   datasets.utils.logging.get_logger(__name__)

_lowerCAmelCase	:   Any    =   ["names", "prefix"]
_lowerCAmelCase	:   str    =   ["warn_bad_lines", "error_bad_lines", "mangle_dupe_cols"]
_lowerCAmelCase	:   List[str]    =   ["encoding_errors", "on_bad_lines"]
_lowerCAmelCase	:   int    =   ["date_format"]


@dataclass
class 	__snake_case   ( datasets.BuilderConfig ):
		SCREAMING_SNAKE_CASE__								=	","
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	"infer"
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	True
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	True
		SCREAMING_SNAKE_CASE__								=	True
		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	True
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	"."
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	'"'
		SCREAMING_SNAKE_CASE__								=	0
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	True
		SCREAMING_SNAKE_CASE__								=	True
		SCREAMING_SNAKE_CASE__								=	0
		SCREAMING_SNAKE_CASE__								=	True
		SCREAMING_SNAKE_CASE__								=	False
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	10000
		SCREAMING_SNAKE_CASE__								=	None
		SCREAMING_SNAKE_CASE__								=	"strict"
		SCREAMING_SNAKE_CASE__								=	"error"
		SCREAMING_SNAKE_CASE__								=	None
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						if self.delimiter is not None:
										lowerCAmelCase__		  =  self.delimiter
						if self.column_names is not None:
										lowerCAmelCase__		  =  self.column_names



		@property
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  {
						    'sep': self.sep,
						    'header': self.header,
						    'names': self.names,
						    'index_col': self.index_col,
						    'usecols': self.usecols,
						    'prefix': self.prefix,
						    'mangle_dupe_cols': self.mangle_dupe_cols,
						    'engine': self.engine,
						    'converters': self.converters,
						    'true_values': self.true_values,
						    'false_values': self.false_values,
						    'skipinitialspace': self.skipinitialspace,
						    'skiprows': self.skiprows,
						    'nrows': self.nrows,
						    'na_values': self.na_values,
						    'keep_default_na': self.keep_default_na,
						    'na_filter': self.na_filter,
						    'verbose': self.verbose,
						    'skip_blank_lines': self.skip_blank_lines,
						    'thousands': self.thousands,
						    'decimal': self.decimal,
						    'lineterminator': self.lineterminator,
						    'quotechar': self.quotechar,
						    'quoting': self.quoting,
						    'escapechar': self.escapechar,
						    'comment': self.comment,
						    'encoding': self.encoding,
						    'dialect': self.dialect,
						    'error_bad_lines': self.error_bad_lines,
						    'warn_bad_lines': self.warn_bad_lines,
						    'skipfooter': self.skipfooter,
						    'doublequote': self.doublequote,
						    'memory_map': self.memory_map,
						    'float_precision': self.float_precision,
						    'chunksize': self.chunksize,
						    'encoding_errors': self.encoding_errors,
						    'on_bad_lines': self.on_bad_lines,
						    'date_format': self.date_format,
						}

						# some kwargs must not be passed if they don't have a default value
						# some others are deprecated and we can also not pass them if they are the default value
						for pd_read_csv_parameter in _PANDAS_READ_CSV_NO_DEFAULT_PARAMETERS + _PANDAS_READ_CSV_DEPRECATED_PARAMETERS:
										if pd_read_csv_kwargs[pd_read_csv_parameter] == getattr(CsvConfig()    ,a_						):
														del pd_read_csv_kwargs[pd_read_csv_parameter]

        # Remove 2.0 new arguments
						if not (datasets.config.PANDAS_VERSION.major >= 2):
										for pd_read_csv_parameter in _PANDAS_READ_CSV_NEW_2_0_0_PARAMETERS:
														del pd_read_csv_kwargs[pd_read_csv_parameter]

        # Remove 1.3 new arguments
						if not (datasets.config.PANDAS_VERSION.major >= 1 and datasets.config.PANDAS_VERSION.minor >= 3):
										for pd_read_csv_parameter in _PANDAS_READ_CSV_NEW_1_3_0_PARAMETERS:
														del pd_read_csv_kwargs[pd_read_csv_parameter]

						return pd_read_csv_kwargs



class 	__snake_case   ( datasets.ArrowBasedBuilder ):
		SCREAMING_SNAKE_CASE__								=	CsvConfig
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return datasets.DatasetInfo(features=self.config.features						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_						):







						"""simple docstring"""
						if not self.config.data_files:
										raise ValueError(f'At least one data file must be specified, but got data_files={self.config.data_files}'						)
						lowerCAmelCase__		  =  dl_manager.download_and_extract(self.config.data_files						)
						if isinstance(a_    ,(str, list, tuple)						):
										lowerCAmelCase__		  =  data_files
										if isinstance(a_    ,a_						):
														lowerCAmelCase__		  =  [files]
										lowerCAmelCase__		  =  [dl_manager.iter_files(a_						) for file in files]
										return [datasets.SplitGenerator(name=datasets.Split.TRAIN    ,gen_kwargs={'files': files}						)]
						lowerCAmelCase__		  =  []
						for split_name, files in data_files.items():
										if isinstance(a_    ,a_						):
														lowerCAmelCase__		  =  [files]
										lowerCAmelCase__		  =  [dl_manager.iter_files(a_						) for file in files]
										splits.append(datasets.SplitGenerator(name=a_    ,gen_kwargs={'files': files}						)						)
						return splits
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_						):







						"""simple docstring"""
						if self.config.features is not None:
										lowerCAmelCase__		  =  self.config.features.arrow_schema
										if all(not require_storage_cast(a_						) for feature in self.config.features.values()						):
														# cheaper cast
														lowerCAmelCase__		  =  pa.Table.from_arrays([pa_table[field.name] for field in schema]    ,schema=a_						)
										else:
														# more expensive cast; allows str <-> int/float or str to Audio for example
														lowerCAmelCase__		  =  table_cast(a_    ,a_						)
						return pa_table



		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.config.features.arrow_schema if self.config.features else None
						# dtype allows reading an int column as str
						lowerCAmelCase__		  =  (
						    {
						        name: dtype.to_pandas_dtype() if not require_storage_cast(a_						) else object
						        for name, dtype, feature in zip(schema.names    ,schema.types    ,self.config.features.values()						)
						    }
						    if schema is not None
						    else None
						)
						for file_idx, file in enumerate(itertools.chain.from_iterable(a_						)						):
										lowerCAmelCase__		  =  pd.read_csv(a_    ,iterator=a_    ,dtype=a_    ,**self.config.pd_read_csv_kwargs						)
										try:
														for batch_idx, df in enumerate(a_						):
																		lowerCAmelCase__		  =  pa.Table.from_pandas(a_						)
																		# Uncomment for debugging (will print the Arrow table size and elements)
																		# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																		# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																		yield (file_idx, batch_idx), self._cast_table(a_						)
										except ValueError as e:
														logger.error(f'Failed to read file \'{file}\' with error {type(a_						)}: {e}'						)
														raise
