codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert switchtransformersx checkpoints from the original repository to jaxflax model import argparse import re from flax traverseutil import flattendict unflattendict from t5x import checkpoints from transformers import switchtransformersconfig switchtransformersforconditionalgeneration from transformers modelingflaxpytorchutils import loadflaxweightsinpytorchmodel from transformers utils import logging logging setverbosityinfo should not include what is already done by the frompt argument moelayernamemapping attention 0selfattention selfattention 0selfattention encoderdecoderattention 1encdecattention value v query q key k out o preselfattentionlayernorm 0layernorm precrossattentionlayernorm 1layernorm preattentionlayernorm 0layernorm previously 1 but seems wrong tokenembedder shared encodernorm finallayernorm decodernorm finallayernorm relposbiasrelembedding block0layer0selfattentionrelativeattentionbiasweight routerrouterweightsw routerclassifier roerroerweightsw routerclassifier logitsdense lmhead def renamekeyssdict 1 in hf t5 we have block x layer y which corresponds to layer x in the original model keys listsdict keys for key in keys layertoblockoflayer r layersd newkey key if re matchlayertoblockoflayer key newkey re subrlayersd rblock1layer newkey layertoblockoflayer rencoderdecoder if re matchlayertoblockoflayer key groups re matchlayertoblockoflayer newkey groups if groups0 encoder newkey re subrmlp r1mlp newkey newkey re subrpremlplayernorm r1layernorm newkey elif groups0 decoder newkey re subrmlp r2mlp newkey newkey re subrpremlplayernorm r2layernorm newkey 2 convert other classic mappings for oldkey tempkey in moelayernamemapping items if oldkey in newkey newkey newkey replaceoldkey tempkey printfkey newkey sdictnewkey sdict popkey if encoderblock0layer0selfattentionrelativeattentionbiasweight in sdict sdictencoderblock0layer0selfattentionrelativeattentionbiasweight sdict encoderblock0layer0selfattentionrelativeattentionbiasweight t if decoderblock0layer0selfattentionrelativeattentionbiasweight in sdict sdictdecoderblock0layer0selfattentionrelativeattentionbiasweight sdict decoderblock0layer0selfattentionrelativeattentionbiasweight t 3 take extra care of the experts layer for key in listsdict keys if expert in key numexperts sdictkey shape0 expertweihts sdictkey for idx in rangenumexperts sdictkey replaceexpert fexpertsexpertidx expertweihtsidx printfkey key replace expert f expertsexpertidx sdict popkey return sdict gintoconfigmapping numencoderlayers numlayers numdecoderlayers numdecoderlayers numheads numheads headdim dkv embeddim dmodel mlpdim dff numselectedexperts numselectedexperts numencodersparselayers numsparseencoderlayers numdecodersparselayers numsparsedecoderlayers dense mlpblock activations feedforwardproj def convertgintoconfigginfile numexperts convert a google style config to the hugging face fromat import regex as re with openginfile r as f rawgin f read regexmatch re findallr 09 rawgin args for param value in regexmatch if param in gintoconfigmapping and value argsgintoconfigmappingparam floatvalue if in value else intvalue activation re findallr activations rawgin0 argsgintoconfigmappingactivation0 stractivation1 argsnumexperts numexperts config switchtransformersconfigargs return config def convertflaxcheckpointtopytorch flaxcheckpointpath configfile ginfilenone pytorchdumppath numexperts8 initialise pytorch model printfloading flax weights from flaxcheckpointpath flaxparams checkpoints loadt5xcheckpointflaxcheckpointpath if ginfile is not none config convertgintoconfigginfile numexperts else config switchtransformersconfig frompretrainedconfigfile ptmodel switchtransformersforconditionalgenerationconfig flaxparams flaxparamstarget flaxparams flattendictflaxparams sep flaxparams renamekeysflaxparams flaxparams unflattendictflaxparams sep load the flax params in the pt model loadflaxweightsinpytorchmodelptmodel flaxparams printfsave pytorch model to pytorchdumppath ptmodel savepretrainedpytorchdumppath if name main parser argparse argumentparser required parameters parser addargument switcht5xcheckpointpath defaultnone typestr requiredtrue help the config json file corresponding to the pretrained switchtransformers model nthis specifies the model architecture if not provided a ginfile has to be provided parser addargument ginfile defaultnone typestr requiredfalse helppath to the gin config file if not provided a configfile has to be passed parser addargument configname defaultnone typestr requiredfalse helpconfig name of switchtransformers model parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model parser addargumentnumexperts default8 typeint requiredfalse helpnumber of experts args parser parseargs convertflaxcheckpointtopytorch args switcht5xcheckpointpath args configname args ginfile args pytorchdumpfolderpath args numexperts coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license convert switchtransformersx checkpoints from the original repository to jax flax model should not include what is already done by the from_pt argument previously 1 but seems wrong 1 in hf t5 we have block x layer y which corresponds to layer x in the original model 2 convert other classic mappings 3 take extra care of the experts layer convert a google style config to the hugging face fromat initialise pytorch model load the flax params in the pt model required parameters
import argparse
import re

from flax.traverse_util import flatten_dict, unflatten_dict
from t5x import checkpoints

from transformers import SwitchTransformersConfig, SwitchTransformersForConditionalGeneration
from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model
from transformers.utils import logging


logging.set_verbosity_info()



MOE_LAYER_NAME_MAPPING = {
    "/attention/": "/0/SelfAttention/",
    "/self_attention/": "/0/SelfAttention/",
    "/encoder_decoder_attention/": "/1/EncDecAttention/",
    "value": "v",
    "query": "q",
    "key": "k",
    "out": "o",
    "pre_self_attention_layer_norm": "0/layer_norm",
    "pre_cross_attention_layer_norm": "1/layer_norm",
    "pre_attention_layer_norm": "0/layer_norm",  
    "token_embedder": "shared",
    "encoder_norm": "final_layer_norm",
    "decoder_norm": "final_layer_norm",
    "relpos_bias/rel_embedding": "block/0/layer/0/SelfAttention/relative_attention_bias/weight",
    "router/router_weights/w/": "router/classifier/",
    "roer/roer_weights/w/": "router/classifier/",
    "logits_dense": "lm_head",
}


def rename_keys(s_dict):
    
    
    keys = list(s_dict.keys())
    for key in keys:
        layer_to_block_of_layer = r".*/layers_(\d+)"
        new_key = key
        if re.match(layer_to_block_of_layer, key):
            new_key = re.sub(r"layers_(\d+)", r"block/\1/layer", new_key)

        layer_to_block_of_layer = r"(encoder|decoder)\/"

        if re.match(layer_to_block_of_layer, key):
            groups = re.match(layer_to_block_of_layer, new_key).groups()
            if groups[0] == "encoder":
                new_key = re.sub(r"/mlp/", r"/1/mlp/", new_key)
                new_key = re.sub(r"/pre_mlp_layer_norm/", r"/1/layer_norm/", new_key)

            elif groups[0] == "decoder":
                new_key = re.sub(r"/mlp/", r"/2/mlp/", new_key)
                new_key = re.sub(r"/pre_mlp_layer_norm/", r"/2/layer_norm/", new_key)

        
        for old_key, temp_key in MOE_LAYER_NAME_MAPPING.items():
            if old_key in new_key:
                new_key = new_key.replace(old_key, temp_key)

        print(f"{key} -> {new_key}")
        s_dict[new_key] = s_dict.pop(key)

    if "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
        s_dict["encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"] = s_dict[
            "encoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
        ].T
    if "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight" in s_dict:
        s_dict["decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"] = s_dict[
            "decoder/block/0/layer/0/SelfAttention/relative_attention_bias/weight"
        ].T

    
    for key in list(s_dict.keys()):
        if "expert" in key:
            num_experts = s_dict[key].shape[0]
            expert_weihts = s_dict[key]
            for idx in range(num_experts):
                s_dict[key.replace("expert/", f"experts/expert_{idx}/")] = expert_weihts[idx]
                print(f"{key} -> {key.replace('expert/', f'experts/expert_{idx}/')}")

            s_dict.pop(key)

    return s_dict


GIN_TO_CONFIG_MAPPING = {
    "NUM_ENCODER_LAYERS": "num_layers",
    "NUM_DECODER_LAYERS": "num_decoder_layers",
    "NUM_HEADS": "num_heads",
    "HEAD_DIM": "d_kv",
    "EMBED_DIM": "d_model",
    "MLP_DIM": "d_ff",
    "NUM_SELECTED_EXPERTS": "num_selected_experts",
    "NUM_ENCODER_SPARSE_LAYERS": "num_sparse_encoder_layers",
    "NUM_DECODER_SPARSE_LAYERS": "num_sparse_decoder_layers",
    "dense.MlpBlock.activations": "feed_forward_proj",
}


def convert_gin_to_config(gin_file, num_experts):
    
    import regex as re

    with open(gin_file, "r") as f:
        raw_gin = f.read()

    regex_match = re.findall(r"(.*) = ([0-9.]*)", raw_gin)
    args = {}
    for param, value in regex_match:
        if param in GIN_TO_CONFIG_MAPPING and value != "":
            args[GIN_TO_CONFIG_MAPPING[param]] = float(value) if "." in value else int(value)

    activation = re.findall(r"(.*activations) = \(\'(.*)\',\)", raw_gin)[0]
    args[GIN_TO_CONFIG_MAPPING[activation[0]]] = str(activation[1])

    args["num_experts"] = num_experts
    config = SwitchTransformersConfig(**args)
    return config


def convert_flax_checkpoint_to_pytorch(
    flax_checkpoint_path, config_file, gin_file=None, pytorch_dump_path="./", num_experts=8
):
    

    print(f"Loading flax weights from : {flax_checkpoint_path}")
    flax_params = checkpoints.load_t5x_checkpoint(flax_checkpoint_path)

    if gin_file is not None:
        config = convert_gin_to_config(gin_file, num_experts)
    else:
        config = SwitchTransformersConfig.from_pretrained(config_file)

    pt_model = SwitchTransformersForConditionalGeneration(config)

    flax_params = flax_params["target"]
    flax_params = flatten_dict(flax_params, sep="/")
    flax_params = rename_keys(flax_params)
    flax_params = unflatten_dict(flax_params, sep="/")

    
    load_flax_weights_in_pytorch_model(pt_model, flax_params)

    print(f"Save PyTorch model to {pytorch_dump_path}")
    pt_model.save_pretrained(pytorch_dump_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--switch_t5x_checkpoint_path",
        default=None,
        type=str,
        required=True,
        help=(
            "The config json file corresponding to the pre-trained SwitchTransformers model. \nThis specifies the"
            " model architecture. If not provided, a `gin_file` has to be provided."
        ),
    )
    parser.add_argument(
        "--gin_file",
        default=None,
        type=str,
        required=False,
        help="Path to the gin config file. If not provided, a `config_file` has to be passed   ",
    )
    parser.add_argument(
        "--config_name", default=None, type=str, required=False, help="Config name of SwitchTransformers model."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output pytorch model."
    )
    parser.add_argument("--num_experts", default=8, type=int, required=False, help="Number of experts")
    args = parser.parse_args()
    convert_flax_checkpoint_to_pytorch(
        args.switch_t5x_checkpoint_path,
        args.config_name,
        args.gin_file,
        args.pytorch_dump_folder_path,
        args.num_experts,
    )