codingutf8 2023 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 xmod checkpoint import argparse from pathlib import path import fairseq import torch from fairseq models xmod import xmodmodel as fairseqxmodmodel from packaging import version from transformers import xmodconfig xmodformaskedlm xmodforsequenceclassification from transformers utils import logging if version parsefairseq version version parse0 12 2 raise exceptionrequires fairseq 0 12 2 if version parsefairseq version version parse2 raise exceptionrequires fairseq v2 logging setverbosityinfo logger logging getloggername sampletext hello world samplelanguage enxx def convertxmodcheckpointtopytorch xmodcheckpointpath str pytorchdumpfolderpath str classificationhead bool datadir pathdatabin xmod fairseqxmodmodel frompretrained modelnameorpathstrpathxmodcheckpointpath parent checkpointfilepathxmodcheckpointpath name namexmodbase archxmodbase taskmultilingualmaskedlm datanameorpathstrdatadir bpesentencepiece sentencepiecemodelstrpathxmodcheckpointpath parent sentencepiece bpe model srcdictstrdatadir dict txt xmod eval disable dropout printxmod xmodsentencoder xmod model encoder sentenceencoder config xmodconfig vocabsizexmodsentencoder embedtokens numembeddings hiddensizexmod cfg model encoderembeddim numhiddenlayersxmod cfg model encoderlayers numattentionheadsxmod cfg model encoderattentionheads intermediatesizexmod cfg model encoderffnembeddim maxpositionembeddings514 typevocabsize1 layernormeps1e5 pytorch default used in fairseq prenormxmod cfg model encodernormalizebefore adapterreductionfactorgetattrxmod cfg model bottleneck 2 adapterlayernormxmod cfg model adapterlayernorm adapterreuselayernormxmod cfg model adapterreuselayernorm lnbeforeadapterxmod cfg model lnbeforeadapter languagesxmod cfg model languages if classificationhead config numlabels xmod model classificationheadsmnli outproj weight shape0 printour xmod config config model xmodforsequenceclassificationconfig if classificationhead else xmodformaskedlmconfig model eval now let s copy all the weights embeddings model roberta embeddings wordembeddings weight xmodsentencoder embedtokens weight model roberta embeddings positionembeddings weight xmodsentencoder embedpositions weight model roberta embeddings tokentypeembeddings weight data torch zeroslike model roberta embeddings tokentypeembeddings weight just zero them out bc xmod doesn t use them model roberta embeddings layernorm weight xmodsentencoder layernormembedding weight model roberta embeddings layernorm bias xmodsentencoder layernormembedding bias for i in rangeconfig numhiddenlayers encoder start of layer layer model roberta encoder layeri xmodlayer xmodsentencoder layersi self attention selfattn layer attention self if not xmodlayer selfattn kproj weight data shape xmodlayer selfattn qproj weight data shape xmodlayer selfattn vproj weight data shape torch sizeconfig hiddensize config hiddensize raise assertionerrordimensions of selfattention weights do not match selfattn query weight data xmodlayer selfattn qproj weight selfattn query bias data xmodlayer selfattn qproj bias selfattn key weight data xmodlayer selfattn kproj weight selfattn key bias data xmodlayer selfattn kproj bias selfattn value weight data xmodlayer selfattn vproj weight selfattn value bias data xmodlayer selfattn vproj bias selfattention output selfoutput layer attention output if selfoutput dense weight shape xmodlayer selfattn outproj weight shape raise assertionerrordimensions of selfattention output weights do not match selfoutput dense weight xmodlayer selfattn outproj weight selfoutput dense bias xmodlayer selfattn outproj bias selfoutput layernorm weight xmodlayer selfattnlayernorm weight selfoutput layernorm bias xmodlayer selfattnlayernorm bias intermediate intermediate layer intermediate if intermediate dense weight shape xmodlayer fc1 weight shape raise assertionerrordimensions of intermediate weights do not match intermediate dense weight xmodlayer fc1 weight intermediate dense bias xmodlayer fc1 bias output bertoutput layer output if bertoutput dense weight shape xmodlayer fc2 weight shape raise assertionerrordimensions of feedforward weights do not match bertoutput dense weight xmodlayer fc2 weight bertoutput dense bias xmodlayer fc2 bias bertoutput layernorm weight xmodlayer finallayernorm weight bertoutput layernorm bias xmodlayer finallayernorm bias if bertoutput adapterlayernorm is not none bertoutput adapterlayernorm weight xmodlayer adapterlayernorm weight bertoutput adapterlayernorm bias xmodlayer adapterlayernorm bias if sortedbertoutput adaptermodules keys sortedxmodlayer adaptermodules keys raise assertionerrorlists of language adapters do not match for langcode adapter in xmodlayer adaptermodules items toadapter bertoutput adaptermoduleslangcode fromadapter xmodlayer adaptermoduleslangcode toadapter dense1 weight fromadapter fc1 weight toadapter dense1 bias fromadapter fc1 bias toadapter dense2 weight fromadapter fc2 weight toadapter dense2 bias fromadapter fc2 bias end of layer if xmodsentencoder layernorm is not none model roberta encoder layernorm weight xmodsentencoder layernorm weight model roberta encoder layernorm bias xmodsentencoder layernorm bias if classificationhead model classifier dense weight xmod model classificationheadsmnli dense weight model classifier dense bias xmod model classificationheadsmnli dense bias model classifier outproj weight xmod model classificationheadsmnli outproj weight model classifier outproj bias xmod model classificationheadsmnli outproj bias else lm head model lmhead dense weight xmod model encoder lmhead dense weight model lmhead dense bias xmod model encoder lmhead dense bias model lmhead layernorm weight xmod model encoder lmhead layernorm weight model lmhead layernorm bias xmod model encoder lmhead layernorm bias model lmhead decoder weight xmod model encoder lmhead weight model lmhead decoder bias xmod model encoder lmhead bias let s check that we get the same results inputids xmod encodesampletext unsqueeze0 batch of size 1 model roberta setdefaultlanguagesamplelanguage ouroutput modelinputids0 if classificationhead theiroutput xmod model classificationheadsmnlixmod extractfeaturesinputids else theiroutput xmod modelinputids langidsamplelanguage0 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 pathpytorchdumpfolderpath mkdirparentstrue existoktrue printfsaving model to pytorchdumpfolderpath model savepretrainedpytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument xmodcheckpointpath 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 convertxmodcheckpointtopytorch args xmodcheckpointpath args pytorchdumpfolderpath args classificationhead coding utf 8 2023 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 x mod checkpoint disable dropout pytorch default used in fairseq now let s copy all the weights embeddings just zero them out b c xmod 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
from pathlib import Path

import fairseq
import torch
from fairseq.models.xmod import XMODModel as FairseqXmodModel
from packaging import version

from transformers import XmodConfig, XmodForMaskedLM, XmodForSequenceClassification
from transformers.utils import logging


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

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

SAMPLE_TEXT = "Hello, World!"
SAMPLE_LANGUAGE = "en_XX"


def convert_xmod_checkpoint_to_pytorch(
    xmod_checkpoint_path: str, pytorch_dump_folder_path: str, classification_head: bool
):
    data_dir = Path("data_bin")
    xmod = FairseqXmodModel.from_pretrained(
        model_name_or_path=str(Path(xmod_checkpoint_path).parent),
        checkpoint_file=Path(xmod_checkpoint_path).name,
        _name="xmod_base",
        arch="xmod_base",
        task="multilingual_masked_lm",
        data_name_or_path=str(data_dir),
        bpe="sentencepiece",
        sentencepiece_model=str(Path(xmod_checkpoint_path).parent / "sentencepiece.bpe.model"),
        src_dict=str(data_dir / "dict.txt"),
    )
    xmod.eval()  
    print(xmod)

    xmod_sent_encoder = xmod.model.encoder.sentence_encoder
    config = XmodConfig(
        vocab_size=xmod_sent_encoder.embed_tokens.num_embeddings,
        hidden_size=xmod.cfg.model.encoder_embed_dim,
        num_hidden_layers=xmod.cfg.model.encoder_layers,
        num_attention_heads=xmod.cfg.model.encoder_attention_heads,
        intermediate_size=xmod.cfg.model.encoder_ffn_embed_dim,
        max_position_embeddings=514,
        type_vocab_size=1,
        layer_norm_eps=1e-5,  
        pre_norm=xmod.cfg.model.encoder_normalize_before,
        adapter_reduction_factor=getattr(xmod.cfg.model, "bottleneck", 2),
        adapter_layer_norm=xmod.cfg.model.adapter_layer_norm,
        adapter_reuse_layer_norm=xmod.cfg.model.adapter_reuse_layer_norm,
        ln_before_adapter=xmod.cfg.model.ln_before_adapter,
        languages=xmod.cfg.model.languages,
    )
    if classification_head:
        config.num_labels = xmod.model.classification_heads["mnli"].out_proj.weight.shape[0]

    print("Our X-MOD config:", config)

    model = XmodForSequenceClassification(config) if classification_head else XmodForMaskedLM(config)
    model.eval()

    
    
    model.roberta.embeddings.word_embeddings.weight = xmod_sent_encoder.embed_tokens.weight
    model.roberta.embeddings.position_embeddings.weight = xmod_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 = xmod_sent_encoder.layernorm_embedding.weight
    model.roberta.embeddings.LayerNorm.bias = xmod_sent_encoder.layernorm_embedding.bias

    for i in range(config.num_hidden_layers):
        
        layer = model.roberta.encoder.layer[i]
        xmod_layer = xmod_sent_encoder.layers[i]

        
        self_attn = layer.attention.self
        if not (
            xmod_layer.self_attn.k_proj.weight.data.shape
            == xmod_layer.self_attn.q_proj.weight.data.shape
            == xmod_layer.self_attn.v_proj.weight.data.shape
            == torch.Size((config.hidden_size, config.hidden_size))
        ):
            raise AssertionError("Dimensions of self-attention weights do not match.")

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

        
        self_output = layer.attention.output
        if self_output.dense.weight.shape != xmod_layer.self_attn.out_proj.weight.shape:
            raise AssertionError("Dimensions of self-attention output weights do not match.")
        self_output.dense.weight = xmod_layer.self_attn.out_proj.weight
        self_output.dense.bias = xmod_layer.self_attn.out_proj.bias
        self_output.LayerNorm.weight = xmod_layer.self_attn_layer_norm.weight
        self_output.LayerNorm.bias = xmod_layer.self_attn_layer_norm.bias

        
        intermediate = layer.intermediate
        if intermediate.dense.weight.shape != xmod_layer.fc1.weight.shape:
            raise AssertionError("Dimensions of intermediate weights do not match.")
        intermediate.dense.weight = xmod_layer.fc1.weight
        intermediate.dense.bias = xmod_layer.fc1.bias

        
        bert_output = layer.output
        if bert_output.dense.weight.shape != xmod_layer.fc2.weight.shape:
            raise AssertionError("Dimensions of feed-forward weights do not match.")
        bert_output.dense.weight = xmod_layer.fc2.weight
        bert_output.dense.bias = xmod_layer.fc2.bias
        bert_output.LayerNorm.weight = xmod_layer.final_layer_norm.weight
        bert_output.LayerNorm.bias = xmod_layer.final_layer_norm.bias
        if bert_output.adapter_layer_norm is not None:
            bert_output.adapter_layer_norm.weight = xmod_layer.adapter_layer_norm.weight
            bert_output.adapter_layer_norm.bias = xmod_layer.adapter_layer_norm.bias

        if sorted(bert_output.adapter_modules.keys()) != sorted(xmod_layer.adapter_modules.keys()):
            raise AssertionError("Lists of language adapters do not match.")
        for lang_code, adapter in xmod_layer.adapter_modules.items():
            to_adapter = bert_output.adapter_modules[lang_code]
            from_adapter = xmod_layer.adapter_modules[lang_code]
            to_adapter.dense1.weight = from_adapter.fc1.weight
            to_adapter.dense1.bias = from_adapter.fc1.bias
            to_adapter.dense2.weight = from_adapter.fc2.weight
            to_adapter.dense2.bias = from_adapter.fc2.bias

        

    if xmod_sent_encoder.layer_norm is not None:
        model.roberta.encoder.LayerNorm.weight = xmod_sent_encoder.layer_norm.weight
        model.roberta.encoder.LayerNorm.bias = xmod_sent_encoder.layer_norm.bias

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

    
    input_ids = xmod.encode(SAMPLE_TEXT).unsqueeze(0)  
    model.roberta.set_default_language(SAMPLE_LANGUAGE)

    our_output = model(input_ids)[0]
    if classification_head:
        their_output = xmod.model.classification_heads["mnli"](xmod.extract_features(input_ids))
    else:
        their_output = xmod.model(input_ids, lang_id=[SAMPLE_LANGUAGE])[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")

    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(
        "--xmod_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_xmod_checkpoint_to_pytorch(
        args.xmod_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
    )