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 data2vec checkpoint import argparse import os import pathlib import fairseq import torch from fairseq modules import transformersentenceencoderlayer from packaging import version from transformers import data2vectextconfig data2vectextformaskedlm data2vectextforsequenceclassification data2vectextmodel from transformers models bert modelingbert import bertintermediate bertlayer bertoutput bertselfattention bertselfoutput important in order for this script to run please make sure to download the dictionary dict txt from wget https dl fbaipublicfiles comfairseqmodelsroberta large tar gz file copied from https github compytorchfairseqblobmainexamplesdata2vecmodelsdata2vectext py 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 convertdata2veccheckpointtopytorch data2veccheckpointpath str pytorchdumpfolderpath str classificationhead bool data2veccheckpointdir data2veccheckpointfilename os path splitdata2veccheckpointpath data2vec data2vectextmodel frompretrained data2veccheckpointdir checkpointfiledata2veccheckpointfilename data2vec eval disable dropout data2vecmodel data2vec models0 data2vecsentencoder data2vecmodel encoder sentenceencoder config data2vectextconfig vocabsizedata2vecsentencoder embedtokens numembeddings hiddensizedata2vecmodel args encoderembeddim numhiddenlayersdata2vecmodel args encoderlayers numattentionheadsdata2vecmodel args encoderattentionheads intermediatesizedata2vecmodel args encoderffnembeddim maxpositionembeddings514 typevocabsize1 layernormeps1e5 pytorch default used in fairseq if classificationhead config numlabels data2vec model classificationheadsmnli outproj weight shape0 printour bert config config model data2vectextforsequenceclassificationconfig if classificationhead else data2vectextformaskedlmconfig model eval now let s copy all the weights embeddings model data2vectext embeddings wordembeddings weight data2vecsentencoder embedtokens weight model data2vectext embeddings positionembeddings weight data2vecsentencoder embedpositions weight model data2vectext embeddings tokentypeembeddings weight data torch zeroslike model data2vectext embeddings tokentypeembeddings weight just zero them out bc data2vec doesn t use them model data2vectext embeddings layernorm weight data2vecsentencoder layernormembedding weight model data2vectext embeddings layernorm bias data2vecsentencoder layernormembedding bias for i in rangeconfig numhiddenlayers encoder start of layer layer bertlayer model data2vectext encoder layeri data2veclayer transformersentenceencoderlayer data2vecsentencoder layersi self attention selfattn bertselfattention layer attention self assert data2veclayer selfattn kproj weight data shape torch size config hiddensize config hiddensize shape for data2veclayer selfattn kproj weight data should be f torch sizeconfig hiddensize config hiddensize assert data2veclayer selfattn qproj weight data shape torch size config hiddensize config hiddensize shape for data2veclayer selfattn qproj weight data should be f torch sizeconfig hiddensize config hiddensize assert data2veclayer selfattn vproj weight data shape torch size config hiddensize config hiddensize shape for data2veclayer selfattn vproj weight data should be f torch sizeconfig hiddensize config hiddensize selfattn query weight data data2veclayer selfattn qproj weight selfattn query bias data data2veclayer selfattn qproj bias selfattn key weight data data2veclayer selfattn kproj weight selfattn key bias data data2veclayer selfattn kproj bias selfattn value weight data data2veclayer selfattn vproj weight selfattn value bias data data2veclayer selfattn vproj bias selfattention output selfoutput bertselfoutput layer attention output assert selfoutput dense weight shape data2veclayer selfattn outproj weight shape fshape for selfoutput dense weight should be data2veclayer selfattn outproj weight shape selfoutput dense weight data2veclayer selfattn outproj weight selfoutput dense bias data2veclayer selfattn outproj bias selfoutput layernorm weight data2veclayer selfattnlayernorm weight selfoutput layernorm bias data2veclayer selfattnlayernorm bias intermediate intermediate bertintermediate layer intermediate assert intermediate dense weight shape data2veclayer fc1 weight shape fshape for intermediate dense weight should be data2veclayer fc1 weight shape intermediate dense weight data2veclayer fc1 weight intermediate dense bias data2veclayer fc1 bias output bertoutput bertoutput layer output assert bertoutput dense weight shape data2veclayer fc2 weight shape fshape for bertoutput dense weight should be data2veclayer fc2 weight shape bertoutput dense weight data2veclayer fc2 weight bertoutput dense bias data2veclayer fc2 bias bertoutput layernorm weight data2veclayer finallayernorm weight bertoutput layernorm bias data2veclayer finallayernorm bias end of layer if classificationhead model classifier dense weight data2vec model classificationheadsmnli dense weight model classifier dense bias data2vec model classificationheadsmnli dense bias model classifier outproj weight data2vec model classificationheadsmnli outproj weight model classifier outproj bias data2vec model classificationheadsmnli outproj bias else lm head model lmhead dense weight data2vecmodel encoder lmhead dense weight model lmhead dense bias data2vecmodel encoder lmhead dense bias model lmhead layernorm weight data2vecmodel encoder lmhead layernorm weight model lmhead layernorm bias data2vecmodel encoder lmhead layernorm bias model lmhead decoder weight data2vecmodel encoder lmhead weight model lmhead decoder bias data2vecmodel encoder lmhead bias let s check that we get the same results inputids torch tensor data2vec encodesampletext unsqueeze0 batch of size 1 ouroutput modelinputids0 if classificationhead theiroutput data2vec model classificationheadsmnlidata2vec extractfeaturesinputids else theiroutput data2vecmodelinputids0 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 checkpointpath 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 convertdata2veccheckpointtopytorch args checkpointpath args pytorchdumpfolderpath args classificationhead 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 data2vec checkpoint important in order for this script to run please make sure to download the dictionary dict txt from wget https dl fbaipublicfiles com fairseq models roberta large tar gz file copied from https github com pytorch fairseq blob main examples data2vec models data2vec_text py copy paste tweak data2vec 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 data2vec 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 os
import pathlib

import fairseq
import torch
from fairseq.modules import TransformerSentenceEncoderLayer
from packaging import version

from transformers import (
    Data2VecTextConfig,
    Data2VecTextForMaskedLM,
    Data2VecTextForSequenceClassification,
    Data2VecTextModel,
)
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_data2vec_checkpoint_to_pytorch(
    data2vec_checkpoint_path: str, pytorch_dump_folder_path: str, classification_head: bool
):
    
    data2vec_checkpoint_dir, data2vec_checkpoint_file_name = os.path.split(data2vec_checkpoint_path)
    data2vec = Data2VecTextModel.from_pretrained(
        data2vec_checkpoint_dir, checkpoint_file=data2vec_checkpoint_file_name
    )
    data2vec.eval()  
    data2vec_model = data2vec.models[0]
    data2vec_sent_encoder = data2vec_model.encoder.sentence_encoder
    config = Data2VecTextConfig(
        vocab_size=data2vec_sent_encoder.embed_tokens.num_embeddings,
        hidden_size=data2vec_model.args.encoder_embed_dim,
        num_hidden_layers=data2vec_model.args.encoder_layers,
        num_attention_heads=data2vec_model.args.encoder_attention_heads,
        intermediate_size=data2vec_model.args.encoder_ffn_embed_dim,
        max_position_embeddings=514,
        type_vocab_size=1,
        layer_norm_eps=1e-5,  
    )
    if classification_head:
        config.num_labels = data2vec.model.classification_heads["mnli"].out_proj.weight.shape[0]
    print("Our BERT config:", config)

    model = Data2VecTextForSequenceClassification(config) if classification_head else Data2VecTextForMaskedLM(config)
    model.eval()

    
    
    model.data2vec_text.embeddings.word_embeddings.weight = data2vec_sent_encoder.embed_tokens.weight
    model.data2vec_text.embeddings.position_embeddings.weight = data2vec_sent_encoder.embed_positions.weight
    model.data2vec_text.embeddings.token_type_embeddings.weight.data = torch.zeros_like(
        model.data2vec_text.embeddings.token_type_embeddings.weight
    )  
    model.data2vec_text.embeddings.LayerNorm.weight = data2vec_sent_encoder.layernorm_embedding.weight
    model.data2vec_text.embeddings.LayerNorm.bias = data2vec_sent_encoder.layernorm_embedding.bias

    for i in range(config.num_hidden_layers):
        
        layer: BertLayer = model.data2vec_text.encoder.layer[i]
        data2vec_layer: TransformerSentenceEncoderLayer = data2vec_sent_encoder.layers[i]

        
        self_attn: BertSelfAttention = layer.attention.self
        assert data2vec_layer.self_attn.k_proj.weight.data.shape == torch.Size(
            (config.hidden_size, config.hidden_size)
        ), (
            "Shape for data2vec_layer.self_attn.k_proj.weight.data should be"
            f" {torch.Size((config.hidden_size, config.hidden_size))}"
        )
        assert data2vec_layer.self_attn.q_proj.weight.data.shape == torch.Size(
            (config.hidden_size, config.hidden_size)
        ), (
            "Shape for data2vec_layer.self_attn.q_proj.weight.data should be"
            f" {torch.Size((config.hidden_size, config.hidden_size))}"
        )
        assert data2vec_layer.self_attn.v_proj.weight.data.shape == torch.Size(
            (config.hidden_size, config.hidden_size)
        ), (
            "Shape for data2vec_layer.self_attn.v_proj.weight.data should be"
            f" {torch.Size((config.hidden_size, config.hidden_size))}"
        )

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

        
        self_output: BertSelfOutput = layer.attention.output
        assert (
            self_output.dense.weight.shape == data2vec_layer.self_attn.out_proj.weight.shape
        ), f"Shape for self_output.dense.weight should be {data2vec_layer.self_attn.out_proj.weight.shape}"
        self_output.dense.weight = data2vec_layer.self_attn.out_proj.weight
        self_output.dense.bias = data2vec_layer.self_attn.out_proj.bias
        self_output.LayerNorm.weight = data2vec_layer.self_attn_layer_norm.weight
        self_output.LayerNorm.bias = data2vec_layer.self_attn_layer_norm.bias

        
        intermediate: BertIntermediate = layer.intermediate
        assert (
            intermediate.dense.weight.shape == data2vec_layer.fc1.weight.shape
        ), f"Shape for intermediate.dense.weight should be {data2vec_layer.fc1.weight.shape}"
        intermediate.dense.weight = data2vec_layer.fc1.weight
        intermediate.dense.bias = data2vec_layer.fc1.bias

        
        bert_output: BertOutput = layer.output
        assert (
            bert_output.dense.weight.shape == data2vec_layer.fc2.weight.shape
        ), f"Shape for bert_output.dense.weight should be {data2vec_layer.fc2.weight.shape}"
        bert_output.dense.weight = data2vec_layer.fc2.weight
        bert_output.dense.bias = data2vec_layer.fc2.bias
        bert_output.LayerNorm.weight = data2vec_layer.final_layer_norm.weight
        bert_output.LayerNorm.bias = data2vec_layer.final_layer_norm.bias
        

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

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

    our_output = model(input_ids)[0]
    if classification_head:
        their_output = data2vec.model.classification_heads["mnli"](data2vec.extract_features(input_ids))
    else:
        their_output = data2vec_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(
        "--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_data2vec_checkpoint_to_pytorch(
        args.checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
    )