codingutf8 2023 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 mra checkpoints from the original repository url https github commlpenmraattention import argparse import torch from transformers import mraconfig mraformaskedlm def renamekeyorigkey if model in origkey origkey origkey replacemodel if norm1 in origkey origkey origkey replacenorm1 attention output layernorm if norm2 in origkey origkey origkey replacenorm2 output layernorm if norm in origkey origkey origkey replacenorm layernorm if transformer in origkey layernum origkey split 0 split1 origkey origkey replaceftransformerlayernum fencoder layer layernum if mha attn in origkey origkey origkey replacemha attn attention self if mha in origkey origkey origkey replacemha attention if wq in origkey origkey origkey replacewq self query if wk in origkey origkey origkey replacewk self key if wv in origkey origkey origkey replacewv self value if ff 0 in origkey origkey origkey replaceff 0 intermediate dense if ff 2 in origkey origkey origkey replaceff 2 output dense if ff in origkey origkey origkey replaceff output dense if mlmclass in origkey origkey origkey replacemlm mlmclass cls predictions decoder if mlm in origkey origkey origkey replacemlm cls predictions transform if backbone backbone encoders in origkey origkey origkey replacebackbone backbone encoders encoder layer if cls not in origkey origkey mra origkey return origkey def convertcheckpointhelpermaxpositionembeddings origstatedict for key in origstatedict copy keys val origstatedict popkey if pooler in key or senclass in key continue else origstatedictrenamekeykey val origstatedictcls predictions bias origstatedictcls predictions decoder bias origstatedictmra embeddings positionids torch arangemaxpositionembeddings expand1 1 2 return origstatedict def convertmracheckpointcheckpointpath mraconfigfile pytorchdumppath origstatedict torch loadcheckpointpath maplocationcpumodelstatedict config mraconfig fromjsonfilemraconfigfile model mraformaskedlmconfig newstatedict convertcheckpointhelperconfig maxpositionembeddings origstatedict printmodel loadstatedictnewstatedict model eval model savepretrainedpytorchdumppath printfcheckpoint successfuly converted model saved at pytorchdumppath if name main parser argparse argumentparser required parameters parser addargument pytorchmodelpath defaultnone typestr requiredtrue helppath to mra pytorch checkpoint parser addargument configfile defaultnone typestr requiredtrue helpthe json file for mra model config parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs convertmracheckpointargs pytorchmodelpath args configfile args pytorchdumppath coding utf 8 2023 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 mra checkpoints from the original repository url https github com mlpen mra attention required parameters
import argparse

import torch

from transformers import MraConfig, MraForMaskedLM


def rename_key(orig_key):
    if "model" in orig_key:
        orig_key = orig_key.replace("model.", "")
    if "norm1" in orig_key:
        orig_key = orig_key.replace("norm1", "attention.output.LayerNorm")
    if "norm2" in orig_key:
        orig_key = orig_key.replace("norm2", "output.LayerNorm")
    if "norm" in orig_key:
        orig_key = orig_key.replace("norm", "LayerNorm")
    if "transformer" in orig_key:
        layer_num = orig_key.split(".")[0].split("_")[-1]
        orig_key = orig_key.replace(f"transformer_{layer_num}", f"encoder.layer.{layer_num}")
    if "mha.attn" in orig_key:
        orig_key = orig_key.replace("mha.attn", "attention.self")
    if "mha" in orig_key:
        orig_key = orig_key.replace("mha", "attention")
    if "W_q" in orig_key:
        orig_key = orig_key.replace("W_q", "self.query")
    if "W_k" in orig_key:
        orig_key = orig_key.replace("W_k", "self.key")
    if "W_v" in orig_key:
        orig_key = orig_key.replace("W_v", "self.value")
    if "ff.0" in orig_key:
        orig_key = orig_key.replace("ff.0", "intermediate.dense")
    if "ff.2" in orig_key:
        orig_key = orig_key.replace("ff.2", "output.dense")
    if "ff" in orig_key:
        orig_key = orig_key.replace("ff", "output.dense")
    if "mlm_class" in orig_key:
        orig_key = orig_key.replace("mlm.mlm_class", "cls.predictions.decoder")
    if "mlm" in orig_key:
        orig_key = orig_key.replace("mlm", "cls.predictions.transform")
    if "backbone.backbone.encoders" in orig_key:
        orig_key = orig_key.replace("backbone.backbone.encoders", "encoder.layer")
    if "cls" not in orig_key:
        orig_key = "mra." + orig_key

    return orig_key


def convert_checkpoint_helper(max_position_embeddings, orig_state_dict):
    for key in orig_state_dict.copy().keys():
        val = orig_state_dict.pop(key)

        if ("pooler" in key) or ("sen_class" in key):
            continue
        else:
            orig_state_dict[rename_key(key)] = val

    orig_state_dict["cls.predictions.bias"] = orig_state_dict["cls.predictions.decoder.bias"]
    orig_state_dict["mra.embeddings.position_ids"] = torch.arange(max_position_embeddings).expand((1, -1)) + 2

    return orig_state_dict


def convert_mra_checkpoint(checkpoint_path, mra_config_file, pytorch_dump_path):
    orig_state_dict = torch.load(checkpoint_path, map_location="cpu")["model_state_dict"]
    config = MraConfig.from_json_file(mra_config_file)
    model = MraForMaskedLM(config)

    new_state_dict = convert_checkpoint_helper(config.max_position_embeddings, orig_state_dict)

    print(model.load_state_dict(new_state_dict))
    model.eval()
    model.save_pretrained(pytorch_dump_path)

    print(f"Checkpoint successfuly converted. Model saved at {pytorch_dump_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--pytorch_model_path", default=None, type=str, required=True, help="Path to Mra pytorch checkpoint."
    )
    parser.add_argument(
        "--config_file",
        default=None,
        type=str,
        required=True,
        help="The json file for Mra model config.",
    )
    parser.add_argument(
        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    args = parser.parse_args()
    convert_mra_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)