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 nystromformer checkpoints from the original repository import argparse import torch from transformers import nystromformerconfig nystromformerformaskedlm 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 ff1 in origkey origkey origkey replaceff1 intermediate dense if ff2 in origkey origkey origkey replaceff2 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 cls not in origkey origkey nystromformer origkey return origkey def convertcheckpointhelperconfig origstatedict for key in origstatedict copy keys val origstatedict popkey if pooler in key or senclass in key or conv bias in key continue else origstatedictrenamekeykey val origstatedictcls predictions bias origstatedictcls predictions decoder bias origstatedictnystromformer embeddings positionids torch arangeconfig maxpositionembeddings expand1 1 2 return origstatedict def convertnystromformercheckpointcheckpointpath nystromformerconfigfile pytorchdumppath origstatedict torch loadcheckpointpath maplocationcpumodelstatedict config nystromformerconfig fromjsonfilenystromformerconfigfile model nystromformerformaskedlmconfig newstatedict convertcheckpointhelperconfig origstatedict model 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 nystromformer pytorch checkpoint parser addargument configfile defaultnone typestr requiredtrue helpthe json file for nystromformer model config parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs convertnystromformercheckpointargs pytorchmodelpath args configfile args pytorchdumppath 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 nystromformer checkpoints from the original repository required parameters
import argparse

import torch

from transformers import NystromformerConfig, NystromformerForMaskedLM


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 "ff1" in orig_key:
        orig_key = orig_key.replace("ff1", "intermediate.dense")
    if "ff2" in orig_key:
        orig_key = orig_key.replace("ff2", "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 "cls" not in orig_key:
        orig_key = "nystromformer." + orig_key

    return orig_key


def convert_checkpoint_helper(config, 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) or ("conv.bias" 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["nystromformer.embeddings.position_ids"] = (
        torch.arange(config.max_position_embeddings).expand((1, -1)) + 2
    )

    return orig_state_dict


def convert_nystromformer_checkpoint(checkpoint_path, nystromformer_config_file, pytorch_dump_path):
    orig_state_dict = torch.load(checkpoint_path, map_location="cpu")["model_state_dict"]
    config = NystromformerConfig.from_json_file(nystromformer_config_file)
    model = NystromformerForMaskedLM(config)

    new_state_dict = convert_checkpoint_helper(config, orig_state_dict)

    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 Nystromformer pytorch checkpoint."
    )
    parser.add_argument(
        "--config_file",
        default=None,
        type=str,
        required=True,
        help="The json file for Nystromformer 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_nystromformer_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)