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

__UpperCAmelCase					=[
    "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",
]

__UpperCAmelCase					=[
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							)      ->       List[str]:
    __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*).*'''							,    UpperCamelCase__							)[1]							)
    layer_number -= 3
    return f"""h.{layer_number}.""" + key






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Union[str, Any]:
    if dtype == torch.bool:
        return 1 / 8
    __lowerCamelCase            =    re.search(r'''[^\d](\d+)$'''							,    str(UpperCamelCase__							)							)
    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 __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)      ->       List[Any]:
    # Construct model
    if bloom_config_file == "":
        __lowerCamelCase            =    BloomConfig()
    else:
        __lowerCamelCase            =    BloomConfig.from_json_file(UpperCamelCase__							)

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

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

        __lowerCamelCase            =    None

        __lowerCamelCase            =    BloomConfig()

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

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

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

                if tensors is None:
                    __lowerCamelCase            =    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            =    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=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            =    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            =    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(UpperCamelCase__							)							).zfill(5							)							)

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

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

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

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

                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(UpperCamelCase__							) 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=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            =    tensors[key] / pretraining_tp

            __lowerCamelCase            =    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            =    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(UpperCamelCase__							,    exist_ok=UpperCamelCase__							)
        __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()							,    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__":
  __UpperCAmelCase					=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",
  )
  __UpperCAmelCase					=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 doctest
import glob
import importlib
import inspect
import os
import re
from contextlib import contextmanager
from functools import wraps
from unittest.mock import patch

import numpy as np
import pytest
from absl.testing import parameterized

import datasets
from datasets import load_metric

from .utils import for_all_test_methods, local, slow


# mark all tests as integration
__UpperCAmelCase					=pytest.mark.integration


__UpperCAmelCase					={"comet"}
__UpperCAmelCase					=importlib.util.find_spec("fairseq") is not None

__UpperCAmelCase					={"code_eval"}
__UpperCAmelCase					=os.name == "nt"

__UpperCAmelCase					={"bertscore", "frugalscore", "perplexity"}
__UpperCAmelCase					=importlib.util.find_spec("transformers") is not None






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Any:
    @wraps(UpperCamelCase__							)
    def wrapper(self							,    UpperCamelCase__							):
        if not _has_fairseq and metric_name in REQUIRE_FAIRSEQ:
            self.skipTest('''"test requires Fairseq"'''							)
        else:
            test_case(self							,    UpperCamelCase__							)

    return wrapper






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Any:
    @wraps(UpperCamelCase__							)
    def wrapper(self							,    UpperCamelCase__							):
        if not _has_transformers and metric_name in REQUIRE_TRANSFORMERS:
            self.skipTest('''"test requires transformers"'''							)
        else:
            test_case(self							,    UpperCamelCase__							)

    return wrapper






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       List[Any]:
    @wraps(UpperCamelCase__							)
    def wrapper(self							,    UpperCamelCase__							):
        if _on_windows and metric_name in UNSUPPORTED_ON_WINDOWS:
            self.skipTest('''"test not supported on Windows"'''							)
        else:
            test_case(self							,    UpperCamelCase__							)

    return wrapper






def __lowerCAmelCase					(				)      ->       Any:
    __lowerCamelCase            =    [metric_dir.split(os.sep							)[-2] for metric_dir in glob.glob('''./metrics/*/'''							)]
    return [{"testcase_name": x, "metric_name": x} for x in metrics if x != "gleu"]  # gleu is unfinished




@parameterized.named_parameters(get_local_metric_names()					)
@for_all_test_methods(
    UpperCAmelCase__				,  UpperCAmelCase__				,  UpperCAmelCase__					)
@local
class        a__ (  parameterized.TestCase					):
     lowerCamelCase    :	str				={}
     lowerCamelCase    :	Union[str, Any]				=None






     @pytest.mark.filterwarnings('''ignore:metric_module_factory is deprecated:FutureWarning'''       )
     @pytest.mark.filterwarnings('''ignore:load_metric is deprecated:FutureWarning'''       )
     def   SCREAMING_SNAKE_CASE__  (					self						:						Tuple    ,							a						:						List[Any]       ):





         """simple docstring"""



         __lowerCamelCase            =    '''[...]'''
         __lowerCamelCase            =    importlib.import_module(
             datasets.load.metric_module_factory(os.path.join('''metrics'''    ,							a       )       ).module_path       )
         __lowerCamelCase            =    datasets.load.import_main_class(metric_module.__name__    ,							dataset=a       )
         # check parameters
         __lowerCamelCase            =    inspect.signature(metric._compute       ).parameters
         self.assertTrue(all(p.kind != p.VAR_KEYWORD for p in parameters.values()       )       )  # no **kwargs
         # run doctest
         with self.patch_intensive_calls(a    ,							metric_module.__name__       ):
             with self.use_local_metrics():
                 try:
                     __lowerCamelCase            =    doctest.testmod(a    ,							verbose=a    ,							raise_on_error=a       )
                 except doctest.UnexpectedException as e:
                     raise e.exc_info[1]  # raise the exception that doctest caught
         self.assertEqual(results.failed    ,							0       )
         self.assertGreater(results.attempted    ,							1       )






     @slow
     def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[Any]    ,							a						:						Optional[int]       ):





         """simple docstring"""



         __lowerCamelCase            =    '''[...]'''
         __lowerCamelCase            =    importlib.import_module(
             datasets.load.metric_module_factory(os.path.join('''metrics'''    ,							a       )       ).module_path       )
         # run doctest
         with self.use_local_metrics():
             __lowerCamelCase            =    doctest.testmod(a    ,							verbose=a    ,							raise_on_error=a       )
         self.assertEqual(results.failed    ,							0       )
         self.assertGreater(results.attempted    ,							1       )






     @contextmanager
     def   SCREAMING_SNAKE_CASE__  (					self						:						Any    ,							a						:						Optional[Any]    ,							a						:						Any       ):





         """simple docstring"""



         if metric_name in self.INTENSIVE_CALLS_PATCHER:
             with self.INTENSIVE_CALLS_PATCHER[metric_name](a       ):
                 yield
         else:
             yield






     @contextmanager
     def   SCREAMING_SNAKE_CASE__  (					self						:						Dict       ):





         """simple docstring"""



         def load_local_metric(a						:						List[str]    ,							*a						:						Optional[int]    ,							**a						:						Tuple       ):
             return load_metric(os.path.join('''metrics'''    ,							a       )    ,							*a    ,							**a       )

         with patch('''datasets.load_metric'''       ) as mock_load_metric:
             __lowerCamelCase            =    load_local_metric
             yield






     @classmethod
     def   SCREAMING_SNAKE_CASE__  (					cls						:						List[str]    ,							a						:						Tuple       ):





         """simple docstring"""



         def wrapper(a						:						List[Any]       ):
             __lowerCamelCase            =    contextmanager(a       )
             __lowerCamelCase            =    patcher
             return patcher

         return wrapper






@LocalMetricTest.register_intensive_calls_patcher('''bleurt'''							)
def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Union[str, Any]:
    import tensorflow.compat.va as tf
    from bleurt.score import Predictor

    tf.flags.DEFINE_string('''sv'''							,    ''''''							,    ''''''							)  # handle pytest cli flags




    class        a__ (  UpperCAmelCase__					):






         def   SCREAMING_SNAKE_CASE__  (					self						:						Tuple    ,							a						:						Dict       ):





             """simple docstring"""



             assert len(input_dict['''input_ids''']       ) == 2
             return np.array([1.03, 1.04]       )

    # mock predict_fn which is supposed to do a forward pass with a bleurt model
    with patch('''bleurt.score._create_predictor'''							) as mock_create_predictor:
        __lowerCamelCase            =    MockedPredictor()
        yield






@LocalMetricTest.register_intensive_calls_patcher('''bertscore'''							)
def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Tuple:
    import torch

    def bert_cos_score_idf(UpperCamelCase__							,    UpperCamelCase__							,    *UpperCamelCase__							,    **UpperCamelCase__							):
        return torch.tensor([[1.0, 1.0, 1.0]] * len(UpperCamelCase__							)							)

    # mock get_model which is supposed to do download a bert model
    # mock bert_cos_score_idf which is supposed to do a forward pass with a bert model
    with patch('''bert_score.scorer.get_model'''							), patch(
        '''bert_score.scorer.bert_cos_score_idf'''							) as mock_bert_cos_score_idf:
        __lowerCamelCase            =    bert_cos_score_idf
        yield






@LocalMetricTest.register_intensive_calls_patcher('''comet'''							)
def __lowerCAmelCase					(				UpperCamelCase__							)      ->       List[str]:
    def load_from_checkpoint(UpperCamelCase__							):




        class        a__ :






             def   SCREAMING_SNAKE_CASE__  (					self						:						Optional[int]    ,							a						:						str    ,							*a						:						int    ,							**a						:						Tuple       ):





                 """simple docstring"""



                 assert len(a       ) == 2
                 __lowerCamelCase            =    [0.19, 0.92]
                 return scores, sum(a       ) / len(a       )
        return Model()

    # mock load_from_checkpoint which is supposed to do download a bert model
    # mock load_from_checkpoint which is supposed to do download a bert model
    with patch('''comet.download_model'''							) as mock_download_model:
        __lowerCamelCase            =    None
        with patch('''comet.load_from_checkpoint'''							) as mock_load_from_checkpoint:
            __lowerCamelCase            =    load_from_checkpoint
            yield






def __lowerCAmelCase					(				)      ->       List[Any]:
    __lowerCamelCase            =    load_metric(os.path.join('''metrics'''							,    '''seqeval'''							)							)
    __lowerCamelCase            =    '''ERROR'''
    __lowerCamelCase            =    f"""Scheme should be one of [IOB1, IOB2, IOE1, IOE2, IOBES, BILOU], got {wrong_scheme}"""
    with pytest.raises(UpperCamelCase__							,    match=re.escape(UpperCamelCase__							)							):
        metric.compute(predictions=[]							,    references=[]							,    scheme=UpperCamelCase__							)
