codingutf8 2020 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 blenderbot checkpoint import argparse import torch from transformers import blenderbotconfig blenderbotforconditionalgeneration from transformers utils import logging logging setverbosityinfo logger logging getloggername patterns attention attn encoderattention encoderattn qlin qproj klin kproj vlin vproj outlin outproj normembeddings layernormembedding positionembeddings embedpositions embeddings embedtokens ffn lin fc def renamestatedictkeyk if k embeddings weight return shared weight for parlainame hfname in patterns k k replaceparlainame hfname if k startswithencoder k k replace attn selfattn k k replacenorm1 selfattnlayernorm k k replacenorm2 finallayernorm elif k startswithdecoder k k replacenorm1 selfattnlayernorm k k replacenorm2 encoderattnlayernorm k k replacenorm3 finallayernorm return k def renamelayernormkeyssd keys model encoder layernormembedding weight model encoder layernormembedding bias model decoder layernormembedding weight model decoder layernormembedding bias for k in keys v sd popk newk k replacelayernormembedding layernorm assert newk not in sd sdnewk v ignorekeys start torch nograd def convertparlaicheckpointcheckpointpath pytorchdumpfolderpath configjsonpath model torch loadcheckpointpath maplocationcpu sd modelmodel cfg blenderbotconfig fromjsonfileconfigjsonpath m blenderbotforconditionalgenerationcfg validkeys m model statedict keys failures mapping for k v in sd items if k in ignorekeys continue newk renamestatedictkeyk if newk not in validkeys failures appendk newk else mappingnewk v if cfg normalizebefore blenderbot3b checkpoints rename layernormembedding layernorm renamelayernormkeyssd m model loadstatedictmapping stricttrue m half m savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargumentsrcpath typestr helplike blenderbotmodel bin parser addargumentsavedir defaulthfblenderbot typestr helpwhere to save converted model parser addargument hfconfigjson defaultblenderbot3bconfig json typestr helppath to config to use args parser parseargs convertparlaicheckpointargs srcpath args savedir args hfconfigjson coding utf 8 2020 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 blenderbot checkpoint copy paste tweak model s weights to our bert structure blenderbot 3b checkpoints rename layernorm_embedding layer_norm required parameters
import argparse

import torch

from transformers import BlenderbotConfig, BlenderbotForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)

PATTERNS = [
    ["attention", "attn"],
    ["encoder_attention", "encoder_attn"],
    ["q_lin", "q_proj"],
    ["k_lin", "k_proj"],
    ["v_lin", "v_proj"],
    ["out_lin", "out_proj"],
    ["norm_embeddings", "layernorm_embedding"],
    ["position_embeddings", "embed_positions"],
    ["embeddings", "embed_tokens"],
    ["ffn.lin", "fc"],
]


def rename_state_dict_key(k):
    if k == "embeddings.weight":
        return "shared.weight"

    for parlai_name, hf_name in PATTERNS:
        k = k.replace(parlai_name, hf_name)

    if k.startswith("encoder"):
        k = k.replace(".attn", ".self_attn")
        k = k.replace("norm1", "self_attn_layer_norm")
        k = k.replace("norm2", "final_layer_norm")
    elif k.startswith("decoder"):
        k = k.replace("norm1", "self_attn_layer_norm")
        k = k.replace("norm2", "encoder_attn_layer_norm")
        k = k.replace("norm3", "final_layer_norm")
    return k


def rename_layernorm_keys(sd):
    keys = [
        "model.encoder.layernorm_embedding.weight",
        "model.encoder.layernorm_embedding.bias",
        "model.decoder.layernorm_embedding.weight",
        "model.decoder.layernorm_embedding.bias",
    ]
    for k in keys:
        v = sd.pop(k)
        new_k = k.replace("layernorm_embedding", "layer_norm")
        assert new_k not in sd
        sd[new_k] = v


IGNORE_KEYS = ["START"]


@torch.no_grad()
def convert_parlai_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_json_path):
    
    model = torch.load(checkpoint_path, map_location="cpu")
    sd = model["model"]
    cfg = BlenderbotConfig.from_json_file(config_json_path)
    m = BlenderbotForConditionalGeneration(cfg)
    valid_keys = m.model.state_dict().keys()
    failures = []
    mapping = {}
    for k, v in sd.items():
        if k in IGNORE_KEYS:
            continue

        new_k = rename_state_dict_key(k)
        if new_k not in valid_keys:
            failures.append([k, new_k])
        else:
            mapping[new_k] = v
    if cfg.normalize_before:  
        rename_layernorm_keys(sd)
    m.model.load_state_dict(mapping, strict=True)
    m.half()
    m.save_pretrained(pytorch_dump_folder_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument("--src_path", type=str, help="like blenderbot-model.bin")
    parser.add_argument("--save_dir", default="hf_blenderbot", type=str, help="Where to save converted model.")
    parser.add_argument(
        "--hf_config_json", default="blenderbot-3b-config.json", type=str, help="Path to config to use"
    )
    args = parser.parse_args()
    convert_parlai_checkpoint(args.src_path, args.save_dir, args.hf_config_json)