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 openai gpt checkpoint import argparse import json import numpy import torch from transformers models xlm tokenizationxlm import vocabfilesnames from transformers utils import configname weightsname logging logging setverbosityinfo def convertxlmcheckpointtopytorchxlmcheckpointpath pytorchdumpfolderpath load checkpoint chkpt torch loadxlmcheckpointpath maplocationcpu statedict chkptmodel we have the base model one level deeper than the original xlm repository twolevelsstatedict for k v in statedict items if predlayer in k twolevelsstatedictk v else twolevelsstatedicttransformer k v config chkptparams config n v for n v in config items if not isinstancev torch floattensor numpy ndarray vocab chkptdicoword2id vocab s w if s find 1 and i 13 else s replace i for s i in vocab items save pytorchmodel pytorchweightsdumppath pytorchdumpfolderpath weightsname pytorchconfigdumppath pytorchdumpfolderpath configname pytorchvocabdumppath pytorchdumpfolderpath vocabfilesnamesvocabfile printfsave pytorch model to pytorchweightsdumppath torch savetwolevelsstatedict pytorchweightsdumppath printfsave configuration file to pytorchconfigdumppath with openpytorchconfigdumppath w encodingutf8 as f f writejson dumpsconfig indent2 n printfsave vocab file to pytorchconfigdumppath with openpytorchvocabdumppath w encodingutf8 as f f writejson dumpsvocab indent2 n if name main parser argparse argumentparser required parameters parser addargument xlmcheckpointpath defaultnone typestr requiredtrue helppath the official pytorch dump parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs convertxlmcheckpointtopytorchargs xlmcheckpointpath args pytorchdumpfolderpath 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 openai gpt checkpoint load checkpoint we have the base model one level deeper than the original xlm repository save pytorch model required parameters
import argparse
import json

import numpy
import torch

from transformers.models.xlm.tokenization_xlm import VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()


def convert_xlm_checkpoint_to_pytorch(xlm_checkpoint_path, pytorch_dump_folder_path):
    
    chkpt = torch.load(xlm_checkpoint_path, map_location="cpu")

    state_dict = chkpt["model"]

    
    two_levels_state_dict = {}
    for k, v in state_dict.items():
        if "pred_layer" in k:
            two_levels_state_dict[k] = v
        else:
            two_levels_state_dict["transformer." + k] = v

    config = chkpt["params"]
    config = {n: v for n, v in config.items() if not isinstance(v, (torch.FloatTensor, numpy.ndarray))}

    vocab = chkpt["dico_word2id"]
    vocab = {s + "</w>" if s.find("@@") == -1 and i > 13 else s.replace("@@", ""): i for s, i in vocab.items()}

    
    pytorch_weights_dump_path = pytorch_dump_folder_path + "/" + WEIGHTS_NAME
    pytorch_config_dump_path = pytorch_dump_folder_path + "/" + CONFIG_NAME
    pytorch_vocab_dump_path = pytorch_dump_folder_path + "/" + VOCAB_FILES_NAMES["vocab_file"]

    print(f"Save PyTorch model to {pytorch_weights_dump_path}")
    torch.save(two_levels_state_dict, pytorch_weights_dump_path)

    print(f"Save configuration file to {pytorch_config_dump_path}")
    with open(pytorch_config_dump_path, "w", encoding="utf-8") as f:
        f.write(json.dumps(config, indent=2) + "\n")

    print(f"Save vocab file to {pytorch_config_dump_path}")
    with open(pytorch_vocab_dump_path, "w", encoding="utf-8") as f:
        f.write(json.dumps(vocab, indent=2) + "\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--xlm_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."
    )
    args = parser.parse_args()
    convert_xlm_checkpoint_to_pytorch(args.xlm_checkpoint_path, args.pytorch_dump_folder_path)