codingutf8 2018 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 roberta checkpoint import argparse import pathlib import fairseq import torch from fairseq models roberta import robertamodel as fairseqrobertamodel from fairseq modules import transformersentenceencoderlayer from packaging import version from transformers import robertaconfig robertaformaskedlm robertaforsequenceclassification from transformers models bert modelingbert import bertintermediate bertlayer bertoutput bertselfattention bertselfoutput from transformers utils import logging if version parsefairseq version version parse0 9 0 raise exceptionrequires fairseq 0 9 0 logging setverbosityinfo logger logging getloggername sampletext hello world cc herlolip def convertrobertacheckpointtopytorch robertacheckpointpath str pytorchdumpfolderpath str classificationhead bool roberta fairseqrobertamodel frompretrainedrobertacheckpointpath roberta eval disable dropout robertasentencoder roberta model encoder sentenceencoder config robertaconfig vocabsizerobertasentencoder embedtokens numembeddings hiddensizeroberta args encoderembeddim numhiddenlayersroberta args encoderlayers numattentionheadsroberta args encoderattentionheads intermediatesizeroberta args encoderffnembeddim maxpositionembeddings514 typevocabsize1 layernormeps1e5 pytorch default used in fairseq if classificationhead config numlabels roberta model classificationheadsmnli outproj weight shape0 printour bert config config model robertaforsequenceclassificationconfig if classificationhead else robertaformaskedlmconfig model eval now let s copy all the weights embeddings model roberta embeddings wordembeddings weight robertasentencoder embedtokens weight model roberta embeddings positionembeddings weight robertasentencoder embedpositions weight model roberta embeddings tokentypeembeddings weight data torch zeroslike model roberta embeddings tokentypeembeddings weight just zero them out bc roberta doesn t use them model roberta embeddings layernorm weight robertasentencoder emblayernorm weight model roberta embeddings layernorm bias robertasentencoder emblayernorm bias for i in rangeconfig numhiddenlayers encoder start of layer layer bertlayer model roberta encoder layeri robertalayer transformersentenceencoderlayer robertasentencoder layersi self attention selfattn bertselfattention layer attention self assert robertalayer selfattn kproj weight data shape robertalayer selfattn qproj weight data shape robertalayer selfattn vproj weight data shape torch sizeconfig hiddensize config hiddensize selfattn query weight data robertalayer selfattn qproj weight selfattn query bias data robertalayer selfattn qproj bias selfattn key weight data robertalayer selfattn kproj weight selfattn key bias data robertalayer selfattn kproj bias selfattn value weight data robertalayer selfattn vproj weight selfattn value bias data robertalayer selfattn vproj bias selfattention output selfoutput bertselfoutput layer attention output assert selfoutput dense weight shape robertalayer selfattn outproj weight shape selfoutput dense weight robertalayer selfattn outproj weight selfoutput dense bias robertalayer selfattn outproj bias selfoutput layernorm weight robertalayer selfattnlayernorm weight selfoutput layernorm bias robertalayer selfattnlayernorm bias intermediate intermediate bertintermediate layer intermediate assert intermediate dense weight shape robertalayer fc1 weight shape intermediate dense weight robertalayer fc1 weight intermediate dense bias robertalayer fc1 bias output bertoutput bertoutput layer output assert bertoutput dense weight shape robertalayer fc2 weight shape bertoutput dense weight robertalayer fc2 weight bertoutput dense bias robertalayer fc2 bias bertoutput layernorm weight robertalayer finallayernorm weight bertoutput layernorm bias robertalayer finallayernorm bias end of layer if classificationhead model classifier dense weight roberta model classificationheadsmnli dense weight model classifier dense bias roberta model classificationheadsmnli dense bias model classifier outproj weight roberta model classificationheadsmnli outproj weight model classifier outproj bias roberta model classificationheadsmnli outproj bias else lm head model lmhead dense weight roberta model encoder lmhead dense weight model lmhead dense bias roberta model encoder lmhead dense bias model lmhead layernorm weight roberta model encoder lmhead layernorm weight model lmhead layernorm bias roberta model encoder lmhead layernorm bias model lmhead decoder weight roberta model encoder lmhead weight model lmhead decoder bias roberta model encoder lmhead bias let s check that we get the same results inputids torch tensor roberta encodesampletext unsqueeze0 batch of size 1 ouroutput modelinputids0 if classificationhead theiroutput roberta model classificationheadsmnliroberta extractfeaturesinputids else theiroutput roberta modelinputids0 printouroutput shape theiroutput shape maxabsolutediff torch maxtorch absouroutput theiroutput item printfmaxabsolutediff maxabsolutediff 1e7 success torch allcloseouroutput theiroutput atol1e3 printdo both models output the same tensors if success else if not success raise exceptionsomething went wrong pathlib pathpytorchdumpfolderpath mkdirparentstrue existoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument robertacheckpointpath defaultnone typestr requiredtrue helppath the official pytorch dump parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model parser addargument classificationhead actionstoretrue helpwhether to convert a final classification head args parser parseargs convertrobertacheckpointtopytorch args robertacheckpointpath args pytorchdumpfolderpath args classificationhead coding utf 8 2018 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 roberta checkpoint copy paste tweak roberta s weights to our bert structure disable dropout pytorch default used in fairseq now let s copy all the weights embeddings just zero them out b c roberta doesn t use them encoder start of layer self attention self attention output intermediate output end of layer lm head let s check that we get the same results batch of size 1 1e 7 required parameters
import argparse
import pathlib

import fairseq
import torch
from fairseq.models.roberta import RobertaModel as FairseqRobertaModel
from fairseq.modules import TransformerSentenceEncoderLayer
from packaging import version

from transformers import RobertaConfig, RobertaForMaskedLM, RobertaForSequenceClassification
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(fairseq.__version__) < version.parse("0.9.0"):
    raise Exception("requires fairseq >= 0.9.0")


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

SAMPLE_TEXT = "Hello world! cécé herlolip"


def convert_roberta_checkpoint_to_pytorch(
    roberta_checkpoint_path: str, pytorch_dump_folder_path: str, classification_head: bool
):
    
    roberta = FairseqRobertaModel.from_pretrained(roberta_checkpoint_path)
    roberta.eval()  
    roberta_sent_encoder = roberta.model.encoder.sentence_encoder
    config = RobertaConfig(
        vocab_size=roberta_sent_encoder.embed_tokens.num_embeddings,
        hidden_size=roberta.args.encoder_embed_dim,
        num_hidden_layers=roberta.args.encoder_layers,
        num_attention_heads=roberta.args.encoder_attention_heads,
        intermediate_size=roberta.args.encoder_ffn_embed_dim,
        max_position_embeddings=514,
        type_vocab_size=1,
        layer_norm_eps=1e-5,  
    )
    if classification_head:
        config.num_labels = roberta.model.classification_heads["mnli"].out_proj.weight.shape[0]
    print("Our BERT config:", config)

    model = RobertaForSequenceClassification(config) if classification_head else RobertaForMaskedLM(config)
    model.eval()

    
    
    model.roberta.embeddings.word_embeddings.weight = roberta_sent_encoder.embed_tokens.weight
    model.roberta.embeddings.position_embeddings.weight = roberta_sent_encoder.embed_positions.weight
    model.roberta.embeddings.token_type_embeddings.weight.data = torch.zeros_like(
        model.roberta.embeddings.token_type_embeddings.weight
    )  
    model.roberta.embeddings.LayerNorm.weight = roberta_sent_encoder.emb_layer_norm.weight
    model.roberta.embeddings.LayerNorm.bias = roberta_sent_encoder.emb_layer_norm.bias

    for i in range(config.num_hidden_layers):
        
        layer: BertLayer = model.roberta.encoder.layer[i]
        roberta_layer: TransformerSentenceEncoderLayer = roberta_sent_encoder.layers[i]

        
        self_attn: BertSelfAttention = layer.attention.self
        assert (
            roberta_layer.self_attn.k_proj.weight.data.shape
            == roberta_layer.self_attn.q_proj.weight.data.shape
            == roberta_layer.self_attn.v_proj.weight.data.shape
            == torch.Size((config.hidden_size, config.hidden_size))
        )

        self_attn.query.weight.data = roberta_layer.self_attn.q_proj.weight
        self_attn.query.bias.data = roberta_layer.self_attn.q_proj.bias
        self_attn.key.weight.data = roberta_layer.self_attn.k_proj.weight
        self_attn.key.bias.data = roberta_layer.self_attn.k_proj.bias
        self_attn.value.weight.data = roberta_layer.self_attn.v_proj.weight
        self_attn.value.bias.data = roberta_layer.self_attn.v_proj.bias

        
        self_output: BertSelfOutput = layer.attention.output
        assert self_output.dense.weight.shape == roberta_layer.self_attn.out_proj.weight.shape
        self_output.dense.weight = roberta_layer.self_attn.out_proj.weight
        self_output.dense.bias = roberta_layer.self_attn.out_proj.bias
        self_output.LayerNorm.weight = roberta_layer.self_attn_layer_norm.weight
        self_output.LayerNorm.bias = roberta_layer.self_attn_layer_norm.bias

        
        intermediate: BertIntermediate = layer.intermediate
        assert intermediate.dense.weight.shape == roberta_layer.fc1.weight.shape
        intermediate.dense.weight = roberta_layer.fc1.weight
        intermediate.dense.bias = roberta_layer.fc1.bias

        
        bert_output: BertOutput = layer.output
        assert bert_output.dense.weight.shape == roberta_layer.fc2.weight.shape
        bert_output.dense.weight = roberta_layer.fc2.weight
        bert_output.dense.bias = roberta_layer.fc2.bias
        bert_output.LayerNorm.weight = roberta_layer.final_layer_norm.weight
        bert_output.LayerNorm.bias = roberta_layer.final_layer_norm.bias
        

    if classification_head:
        model.classifier.dense.weight = roberta.model.classification_heads["mnli"].dense.weight
        model.classifier.dense.bias = roberta.model.classification_heads["mnli"].dense.bias
        model.classifier.out_proj.weight = roberta.model.classification_heads["mnli"].out_proj.weight
        model.classifier.out_proj.bias = roberta.model.classification_heads["mnli"].out_proj.bias
    else:
        
        model.lm_head.dense.weight = roberta.model.encoder.lm_head.dense.weight
        model.lm_head.dense.bias = roberta.model.encoder.lm_head.dense.bias
        model.lm_head.layer_norm.weight = roberta.model.encoder.lm_head.layer_norm.weight
        model.lm_head.layer_norm.bias = roberta.model.encoder.lm_head.layer_norm.bias
        model.lm_head.decoder.weight = roberta.model.encoder.lm_head.weight
        model.lm_head.decoder.bias = roberta.model.encoder.lm_head.bias

    
    input_ids: torch.Tensor = roberta.encode(SAMPLE_TEXT).unsqueeze(0)  

    our_output = model(input_ids)[0]
    if classification_head:
        their_output = roberta.model.classification_heads["mnli"](roberta.extract_features(input_ids))
    else:
        their_output = roberta.model(input_ids)[0]
    print(our_output.shape, their_output.shape)
    max_absolute_diff = torch.max(torch.abs(our_output - their_output)).item()
    print(f"max_absolute_diff = {max_absolute_diff}")  
    success = torch.allclose(our_output, their_output, atol=1e-3)
    print("Do both models output the same tensors?", "🔥" if success else "💩")
    if not success:
        raise Exception("Something went wRoNg")

    pathlib.Path(pytorch_dump_folder_path).mkdir(parents=True, exist_ok=True)
    print(f"Saving model to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--roberta_checkpoint_path", default=None, type=str, required=True, help="Path the official PyTorch dump."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--classification_head", action="store_true", help="Whether to convert a final classification head."
    )
    args = parser.parse_args()
    convert_roberta_checkpoint_to_pytorch(
        args.roberta_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
    )