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 bort checkpoint import argparse import os import gluonnlp as nlp import mxnet as mx import numpy as np import torch from gluonnlp base import gethomedir from gluonnlp model bert import bertencoder from gluonnlp model utils import loadvocab from gluonnlp vocab import vocab from packaging import version from torch import nn from transformers import bertconfig bertformaskedlm bertmodel robertatokenizer from transformers models bert modelingbert import bertintermediate bertlayer bertoutput bertselfattention bertselfoutput from transformers utils import logging if version parsenlp version version parse0 8 3 raise exceptionrequires gluonnlp 0 8 3 if version parsemx version version parse1 5 0 raise exceptionrequires mxnet 1 5 0 logging setverbosityinfo logger logging getloggername sampletext the nymphenburg palace is a beautiful palace in munich def convertbortcheckpointtopytorchbortcheckpointpath str pytorchdumpfolderpath str original bort configuration bort487681024hparams attentioncell multihead numlayers 4 units 1024 hiddensize 768 maxlength 512 numheads 8 scaled true dropout 0 1 useresidual true embedsize 1024 embeddropout 0 1 wordembed none layernormeps 1e5 tokentypevocabsize 2 predefinedargs bort487681024hparams let s construct the original bort model here taken from official bert implementation see https github comalexabortblobmasterbortbort py encoder bertencoder attentioncellpredefinedargsattentioncell numlayerspredefinedargsnumlayers unitspredefinedargsunits hiddensizepredefinedargshiddensize maxlengthpredefinedargsmaxlength numheadspredefinedargsnumheads scaledpredefinedargsscaled dropoutpredefinedargsdropout outputattentionfalse outputallencodingsfalse useresidualpredefinedargsuseresidual activationpredefinedargs getactivation gelu layernormepspredefinedargs getlayernormeps none vocab information needs to be fetched first it s the same as roberta so robertatokenizer can be used later vocabname openwebtextccnewsstoriesbookscased specify download folder to gluonnlp s vocab gluoncachedir os path joingethomedir models bortvocab loadvocabvocabname none gluoncachedir clsvocab originalbort nlp model bertmodel encoder lenbortvocab unitspredefinedargsunits embedsizepredefinedargsembedsize embeddropoutpredefinedargsembeddropout wordembedpredefinedargswordembed usepoolerfalse usetokentypeembedfalse tokentypevocabsizepredefinedargstokentypevocabsize useclassifierfalse usedecoderfalse originalbort loadparametersbortcheckpointpath castdtypetrue ignoreextratrue params originalbort collectparamswithprefix build our config hfbortconfigjson architectures bertformaskedlm attentionprobsdropoutprob predefinedargsdropout hiddenact gelu hiddendropoutprob predefinedargsdropout hiddensize predefinedargsembedsize initializerrange 0 02 intermediatesize predefinedargshiddensize layernormeps predefinedargslayernormeps maxpositionembeddings predefinedargsmaxlength modeltype bort numattentionheads predefinedargsnumheads numhiddenlayers predefinedargsnumlayers padtokenid 1 2 bert 1 roberta typevocabsize 1 2 bert 1 roberta vocabsize lenbortvocab hfbortconfig bertconfig fromdicthfbortconfigjson hfbortmodel bertformaskedlmhfbortconfig hfbortmodel eval parameter mapping table gluonnlp to transformers denotes layer index gluon parameter transformers parameter encoder layernorm beta bert embeddings layernorm bias encoder layernorm gamma bert embeddings layernorm weight encoder positionweight bert embeddings positionembeddings weight wordembed 0 weight bert embeddings wordembeddings weight encoder transformercells attentioncell projkey bias bert encoder layer attention self key bias encoder transformercells attentioncell projkey weight bert encoder layer attention self key weight encoder transformercells attentioncell projquery bias bert encoder layer attention self query bias encoder transformercells attentioncell projquery weight bert encoder layer attention self query weight encoder transformercells attentioncell projvalue bias bert encoder layer attention self value bias encoder transformercells attentioncell projvalue weight bert encoder layer attention self value weight encoder transformercells ffn ffn2 bias bert encoder layer attention output dense bias encoder transformercells ffn ffn2 weight bert encoder layer attention output dense weight encoder transformercells layernorm beta bert encoder layer attention output layernorm bias encoder transformercells layernorm gamma bert encoder layer attention output layernorm weight encoder transformercells ffn ffn1 bias bert encoder layer intermediate dense bias encoder transformercells ffn ffn1 weight bert encoder layer intermediate dense weight encoder transformercells ffn layernorm beta bert encoder layer output layernorm bias encoder transformercells ffn layernorm gamma bert encoder layer output layernorm weight encoder transformercells proj bias bert encoder layer output dense bias encoder transformercells proj weight bert encoder layer output dense weight helper function to convert mxnet arrays to pytorch def totorchmxarray nn parameter return nn parametertorch floattensormxarray data asnumpy check param shapes and map new hf param back def checkandmapparamshfparam gluonparam shapehf hfparam shape gluonparam totorchparamsgluonparam shapegluon gluonparam shape assert shapehf shapegluon fthe gluon parameter gluonparam has shape shapegluon but expects shape shapehf for transformers return gluonparam hfbortmodel bert embeddings wordembeddings weight checkandmapparams hfbortmodel bert embeddings wordembeddings weight wordembed 0 weight hfbortmodel bert embeddings positionembeddings weight checkandmapparams hfbortmodel bert embeddings positionembeddings weight encoder positionweight hfbortmodel bert embeddings layernorm bias checkandmapparams hfbortmodel bert embeddings layernorm bias encoder layernorm beta hfbortmodel bert embeddings layernorm weight checkandmapparams hfbortmodel bert embeddings layernorm weight encoder layernorm gamma inspired by roberta conversion script we just zero them out bort does not use them hfbortmodel bert embeddings tokentypeembeddings weight data torch zeroslike hfbortmodel bert embeddings tokentypeembeddings weight data for i in rangehfbortconfig numhiddenlayers layer bertlayer hfbortmodel bert encoder layeri self attention selfattn bertselfattention layer attention self selfattn key bias data checkandmapparams selfattn key bias data fencoder transformercells i attentioncell projkey bias selfattn key weight data checkandmapparams selfattn key weight data fencoder transformercells i attentioncell projkey weight selfattn query bias data checkandmapparams selfattn query bias data fencoder transformercells i attentioncell projquery bias selfattn query weight data checkandmapparams selfattn query weight data fencoder transformercells i attentioncell projquery weight selfattn value bias data checkandmapparams selfattn value bias data fencoder transformercells i attentioncell projvalue bias selfattn value weight data checkandmapparams selfattn value weight data fencoder transformercells i attentioncell projvalue weight self attention output selfoutput bertselfoutput layer attention output selfoutput dense bias checkandmapparams selfoutput dense bias fencoder transformercells i proj bias selfoutput dense weight checkandmapparams selfoutput dense weight fencoder transformercells i proj weight selfoutput layernorm bias checkandmapparams selfoutput layernorm bias fencoder transformercells i layernorm beta selfoutput layernorm weight checkandmapparams selfoutput layernorm weight fencoder transformercells i layernorm gamma intermediate intermediate bertintermediate layer intermediate intermediate dense bias checkandmapparams intermediate dense bias fencoder transformercells i ffn ffn1 bias intermediate dense weight checkandmapparams intermediate dense weight fencoder transformercells i ffn ffn1 weight output bertoutput bertoutput layer output bertoutput dense bias checkandmapparams bertoutput dense bias fencoder transformercells i ffn ffn2 bias bertoutput dense weight checkandmapparams bertoutput dense weight fencoder transformercells i ffn ffn2 weight bertoutput layernorm bias checkandmapparams bertoutput layernorm bias fencoder transformercells i ffn layernorm beta bertoutput layernorm weight checkandmapparams bertoutput layernorm weight fencoder transformercells i ffn layernorm gamma save space and energy hfbortmodel half compare output of both models tokenizer robertatokenizer frompretrainedrobertabase inputids tokenizer encodeplussampletextinputids get gluon output gluoninputids mx nd arrayinputids outputgluon originalbortinputsgluoninputids tokentypes get transformer output save and reload model again hfbortmodel savepretrainedpytorchdumpfolderpath hfbortmodel bertmodel frompretrainedpytorchdumpfolderpath hfbortmodel eval inputids tokenizer encodeplussampletext returntensorspt outputhf hfbortmodelinputids0 gluonlayer outputgluon0 asnumpy hflayer outputhf0 detach numpy maxabsolutediff np maxnp abshflayer gluonlayer item success np allclosegluonlayer hflayer atol1e3 if success print both model do output the same tensors else print both model do not output the same tensors printabsolute difference is maxabsolutediff if name main parser argparse argumentparser required parameters parser addargument bortcheckpointpath defaultnone typestr requiredtrue helppath the official bort params file parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs convertbortcheckpointtopytorchargs bortcheckpointpath args pytorchdumpfolderpath 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 bort checkpoint convert the original bort checkpoint based on mxnet and gluonnlp to our bert structure original bort configuration let s construct the original bort model here taken from official bert implementation see https github com alexa bort blob master bort bort py vocab information needs to be fetched first it s the same as roberta so robertatokenizer can be used later specify download folder to gluonnlp s vocab build our config 2 bert 1 roberta 2 bert 1 roberta parameter mapping table gluonnlp to transformers denotes layer index gluon parameter transformers parameter encoder layer_norm beta bert embeddings layernorm bias encoder layer_norm gamma bert embeddings layernorm weight encoder position_weight bert embeddings position_embeddings weight word_embed 0 weight bert embeddings word_embeddings weight encoder transformer_cells attention_cell proj_key bias bert encoder layer attention self key bias encoder transformer_cells attention_cell proj_key weight bert encoder layer attention self key weight encoder transformer_cells attention_cell proj_query bias bert encoder layer attention self query bias encoder transformer_cells attention_cell proj_query weight bert encoder layer attention self query weight encoder transformer_cells attention_cell proj_value bias bert encoder layer attention self value bias encoder transformer_cells attention_cell proj_value weight bert encoder layer attention self value weight encoder transformer_cells ffn ffn_2 bias bert encoder layer attention output dense bias encoder transformer_cells ffn ffn_2 weight bert encoder layer attention output dense weight encoder transformer_cells layer_norm beta bert encoder layer attention output layernorm bias encoder transformer_cells layer_norm gamma bert encoder layer attention output layernorm weight encoder transformer_cells ffn ffn_1 bias bert encoder layer intermediate dense bias encoder transformer_cells ffn ffn_1 weight bert encoder layer intermediate dense weight encoder transformer_cells ffn layer_norm beta bert encoder layer output layernorm bias encoder transformer_cells ffn layer_norm gamma bert encoder layer output layernorm weight encoder transformer_cells proj bias bert encoder layer output dense bias encoder transformer_cells proj weight bert encoder layer output dense weight helper function to convert mxnet arrays to pytorch check param shapes and map new hf param back inspired by roberta conversion script we just zero them out bort does not use them self attention self attention output intermediate output save space and energy compare output of both models get gluon output get transformer output save and reload model again required parameters
import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse("0.8.3"):
    raise Exception("requires gluonnlp == 0.8.3")

if version.parse(mx.__version__) != version.parse("1.5.0"):
    raise Exception("requires mxnet == 1.5.0")

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

SAMPLE_TEXT = "The Nymphenburg Palace is a beautiful palace in Munich!"


def convert_bort_checkpoint_to_pytorch(bort_checkpoint_path: str, pytorch_dump_folder_path: str):
    

    
    bort_4_8_768_1024_hparams = {
        "attention_cell": "multi_head",
        "num_layers": 4,
        "units": 1024,
        "hidden_size": 768,
        "max_length": 512,
        "num_heads": 8,
        "scaled": True,
        "dropout": 0.1,
        "use_residual": True,
        "embed_size": 1024,
        "embed_dropout": 0.1,
        "word_embed": None,
        "layer_norm_eps": 1e-5,
        "token_type_vocab_size": 2,
    }

    predefined_args = bort_4_8_768_1024_hparams

    
    
    
    encoder = BERTEncoder(
        attention_cell=predefined_args["attention_cell"],
        num_layers=predefined_args["num_layers"],
        units=predefined_args["units"],
        hidden_size=predefined_args["hidden_size"],
        max_length=predefined_args["max_length"],
        num_heads=predefined_args["num_heads"],
        scaled=predefined_args["scaled"],
        dropout=predefined_args["dropout"],
        output_attention=False,
        output_all_encodings=False,
        use_residual=predefined_args["use_residual"],
        activation=predefined_args.get("activation", "gelu"),
        layer_norm_eps=predefined_args.get("layer_norm_eps", None),
    )

    
    
    vocab_name = "openwebtext_ccnews_stories_books_cased"

    
    gluon_cache_dir = os.path.join(get_home_dir(), "models")
    bort_vocab = _load_vocab(vocab_name, None, gluon_cache_dir, cls=Vocab)

    original_bort = nlp.model.BERTModel(
        encoder,
        len(bort_vocab),
        units=predefined_args["units"],
        embed_size=predefined_args["embed_size"],
        embed_dropout=predefined_args["embed_dropout"],
        word_embed=predefined_args["word_embed"],
        use_pooler=False,
        use_token_type_embed=False,
        token_type_vocab_size=predefined_args["token_type_vocab_size"],
        use_classifier=False,
        use_decoder=False,
    )

    original_bort.load_parameters(bort_checkpoint_path, cast_dtype=True, ignore_extra=True)
    params = original_bort._collect_params_with_prefix()

    
    hf_bort_config_json = {
        "architectures": ["BertForMaskedLM"],
        "attention_probs_dropout_prob": predefined_args["dropout"],
        "hidden_act": "gelu",
        "hidden_dropout_prob": predefined_args["dropout"],
        "hidden_size": predefined_args["embed_size"],
        "initializer_range": 0.02,
        "intermediate_size": predefined_args["hidden_size"],
        "layer_norm_eps": predefined_args["layer_norm_eps"],
        "max_position_embeddings": predefined_args["max_length"],
        "model_type": "bort",
        "num_attention_heads": predefined_args["num_heads"],
        "num_hidden_layers": predefined_args["num_layers"],
        "pad_token_id": 1,  
        "type_vocab_size": 1,  
        "vocab_size": len(bort_vocab),
    }

    hf_bort_config = BertConfig.from_dict(hf_bort_config_json)
    hf_bort_model = BertForMaskedLM(hf_bort_config)
    hf_bort_model.eval()

    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    

    
    def to_torch(mx_array) -> nn.Parameter:
        return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()))

    
    def check_and_map_params(hf_param, gluon_param):
        shape_hf = hf_param.shape

        gluon_param = to_torch(params[gluon_param])
        shape_gluon = gluon_param.shape

        assert (
            shape_hf == shape_gluon
        ), f"The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers"

        return gluon_param

    hf_bort_model.bert.embeddings.word_embeddings.weight = check_and_map_params(
        hf_bort_model.bert.embeddings.word_embeddings.weight, "word_embed.0.weight"
    )
    hf_bort_model.bert.embeddings.position_embeddings.weight = check_and_map_params(
        hf_bort_model.bert.embeddings.position_embeddings.weight, "encoder.position_weight"
    )
    hf_bort_model.bert.embeddings.LayerNorm.bias = check_and_map_params(
        hf_bort_model.bert.embeddings.LayerNorm.bias, "encoder.layer_norm.beta"
    )
    hf_bort_model.bert.embeddings.LayerNorm.weight = check_and_map_params(
        hf_bort_model.bert.embeddings.LayerNorm.weight, "encoder.layer_norm.gamma"
    )

    
    hf_bort_model.bert.embeddings.token_type_embeddings.weight.data = torch.zeros_like(
        hf_bort_model.bert.embeddings.token_type_embeddings.weight.data
    )

    for i in range(hf_bort_config.num_hidden_layers):
        layer: BertLayer = hf_bort_model.bert.encoder.layer[i]

        
        self_attn: BertSelfAttention = layer.attention.self

        self_attn.key.bias.data = check_and_map_params(
            self_attn.key.bias.data, f"encoder.transformer_cells.{i}.attention_cell.proj_key.bias"
        )

        self_attn.key.weight.data = check_and_map_params(
            self_attn.key.weight.data, f"encoder.transformer_cells.{i}.attention_cell.proj_key.weight"
        )
        self_attn.query.bias.data = check_and_map_params(
            self_attn.query.bias.data, f"encoder.transformer_cells.{i}.attention_cell.proj_query.bias"
        )
        self_attn.query.weight.data = check_and_map_params(
            self_attn.query.weight.data, f"encoder.transformer_cells.{i}.attention_cell.proj_query.weight"
        )
        self_attn.value.bias.data = check_and_map_params(
            self_attn.value.bias.data, f"encoder.transformer_cells.{i}.attention_cell.proj_value.bias"
        )
        self_attn.value.weight.data = check_and_map_params(
            self_attn.value.weight.data, f"encoder.transformer_cells.{i}.attention_cell.proj_value.weight"
        )

        
        self_output: BertSelfOutput = layer.attention.output

        self_output.dense.bias = check_and_map_params(
            self_output.dense.bias, f"encoder.transformer_cells.{i}.proj.bias"
        )
        self_output.dense.weight = check_and_map_params(
            self_output.dense.weight, f"encoder.transformer_cells.{i}.proj.weight"
        )
        self_output.LayerNorm.bias = check_and_map_params(
            self_output.LayerNorm.bias, f"encoder.transformer_cells.{i}.layer_norm.beta"
        )
        self_output.LayerNorm.weight = check_and_map_params(
            self_output.LayerNorm.weight, f"encoder.transformer_cells.{i}.layer_norm.gamma"
        )

        
        intermediate: BertIntermediate = layer.intermediate

        intermediate.dense.bias = check_and_map_params(
            intermediate.dense.bias, f"encoder.transformer_cells.{i}.ffn.ffn_1.bias"
        )
        intermediate.dense.weight = check_and_map_params(
            intermediate.dense.weight, f"encoder.transformer_cells.{i}.ffn.ffn_1.weight"
        )

        
        bert_output: BertOutput = layer.output

        bert_output.dense.bias = check_and_map_params(
            bert_output.dense.bias, f"encoder.transformer_cells.{i}.ffn.ffn_2.bias"
        )
        bert_output.dense.weight = check_and_map_params(
            bert_output.dense.weight, f"encoder.transformer_cells.{i}.ffn.ffn_2.weight"
        )
        bert_output.LayerNorm.bias = check_and_map_params(
            bert_output.LayerNorm.bias, f"encoder.transformer_cells.{i}.ffn.layer_norm.beta"
        )
        bert_output.LayerNorm.weight = check_and_map_params(
            bert_output.LayerNorm.weight, f"encoder.transformer_cells.{i}.ffn.layer_norm.gamma"
        )

    
    hf_bort_model.half()

    
    tokenizer = RobertaTokenizer.from_pretrained("roberta-base")

    input_ids = tokenizer.encode_plus(SAMPLE_TEXT)["input_ids"]

    
    gluon_input_ids = mx.nd.array([input_ids])
    output_gluon = original_bort(inputs=gluon_input_ids, token_types=[])

    
    hf_bort_model.save_pretrained(pytorch_dump_folder_path)
    hf_bort_model = BertModel.from_pretrained(pytorch_dump_folder_path)
    hf_bort_model.eval()

    input_ids = tokenizer.encode_plus(SAMPLE_TEXT, return_tensors="pt")
    output_hf = hf_bort_model(**input_ids)[0]

    gluon_layer = output_gluon[0].asnumpy()
    hf_layer = output_hf[0].detach().numpy()

    max_absolute_diff = np.max(np.abs(hf_layer - gluon_layer)).item()
    success = np.allclose(gluon_layer, hf_layer, atol=1e-3)

    if success:
        print("✔️ Both model do output the same tensors")
    else:
        print("❌ Both model do **NOT** output the same tensors")
        print("Absolute difference is:", max_absolute_diff)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--bort_checkpoint_path", default=None, type=str, required=True, help="Path the official Bort params file."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    args = parser.parse_args()
    convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)