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 transformer xl checkpoint and datasets import argparse import os import pickle import sys import torch from transformers import transfoxlconfig transfoxllmheadmodel loadtfweightsintransfoxl from transformers models deprecated transfoxl import tokenizationtransfoxl as datautils from transformers models deprecated transfoxl tokenizationtransfoxl import corpusname vocabfilesnames from transformers utils import configname weightsname logging logging setverbosityinfo we do this to be able to load python 2 datasets pickles see e g https stackoverflow comquestions2121874pythonpicklingafterchangingamodulesdirectory21219182121918 datautils vocab datautils transfoxltokenizer datautils corpus datautils transfoxlcorpus sys modulesdatautils datautils sys modulesvocabulary datautils def converttransfoxlcheckpointtopytorch tfcheckpointpath transfoxlconfigfile pytorchdumpfolderpath transfoxldatasetfile if transfoxldatasetfile convert a preprocessed corpus see original tensorflow repo with opentransfoxldatasetfile rb as fp corpus pickle loadfp encodinglatin1 save vocabulary and dataset cache as dictionaries should be better than pickles for the longterm pytorchvocabdumppath pytorchdumpfolderpath vocabfilesnamespretrainedvocabfile printfsave vocabulary to pytorchvocabdumppath corpusvocabdict corpus vocab dict torch savecorpusvocabdict pytorchvocabdumppath corpusdictnovocab corpus dict corpusdictnovocab popvocab none pytorchdatasetdumppath pytorchdumpfolderpath corpusname printfsave dataset to pytorchdatasetdumppath torch savecorpusdictnovocab pytorchdatasetdumppath if tfcheckpointpath convert a pretrained tensorflow model configpath os path abspathtransfoxlconfigfile tfpath os path abspathtfcheckpointpath printfconverting transformer xl checkpoint from tfpath with config at configpath initialise pytorch model if transfoxlconfigfile config transfoxlconfig else config transfoxlconfig fromjsonfiletransfoxlconfigfile printfbuilding pytorch model from configuration config model transfoxllmheadmodelconfig model loadtfweightsintransfoxlmodel config tfpath save pytorchmodel pytorchweightsdumppath os path joinpytorchdumpfolderpath weightsname pytorchconfigdumppath os path joinpytorchdumpfolderpath configname printfsave pytorch model to os path abspathpytorchweightsdumppath torch savemodel statedict pytorchweightsdumppath printfsave configuration file to os path abspathpytorchconfigdumppath with openpytorchconfigdumppath w encodingutf8 as f f writeconfig tojsonstring if name main parser argparse argumentparser parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the folder to store the pytorch model or datasetvocab parser addargument tfcheckpointpath default typestr helpan optional path to a tensorflow checkpoint path to be converted parser addargument transfoxlconfigfile default typestr help an optional config json file corresponding to the pretrained bert model n this specifies the model architecture parser addargument transfoxldatasetfile default typestr helpan optional dataset file to be converted in a vocabulary args parser parseargs converttransfoxlcheckpointtopytorch args tfcheckpointpath args transfoxlconfigfile args pytorchdumpfolderpath args transfoxldatasetfile 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 transformer xl checkpoint and datasets we do this to be able to load python 2 datasets pickles see e g https stackoverflow com questions 2121874 python pickling after changing a modules directory 2121918 2121918 convert a pre processed corpus see original tensorflow repo save vocabulary and dataset cache as dictionaries should be better than pickles for the long term convert a pre trained tensorflow model initialise pytorch model save pytorch model
import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.deprecated.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.deprecated.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()



data_utils.Vocab = data_utils.TransfoXLTokenizer
data_utils.Corpus = data_utils.TransfoXLCorpus
sys.modules["data_utils"] = data_utils
sys.modules["vocabulary"] = data_utils


def convert_transfo_xl_checkpoint_to_pytorch(
    tf_checkpoint_path, transfo_xl_config_file, pytorch_dump_folder_path, transfo_xl_dataset_file
):
    if transfo_xl_dataset_file:
        
        with open(transfo_xl_dataset_file, "rb") as fp:
            corpus = pickle.load(fp, encoding="latin1")
        
        pytorch_vocab_dump_path = pytorch_dump_folder_path + "/" + VOCAB_FILES_NAMES["pretrained_vocab_file"]
        print(f"Save vocabulary to {pytorch_vocab_dump_path}")
        corpus_vocab_dict = corpus.vocab.__dict__
        torch.save(corpus_vocab_dict, pytorch_vocab_dump_path)

        corpus_dict_no_vocab = corpus.__dict__
        corpus_dict_no_vocab.pop("vocab", None)
        pytorch_dataset_dump_path = pytorch_dump_folder_path + "/" + CORPUS_NAME
        print(f"Save dataset to {pytorch_dataset_dump_path}")
        torch.save(corpus_dict_no_vocab, pytorch_dataset_dump_path)

    if tf_checkpoint_path:
        
        config_path = os.path.abspath(transfo_xl_config_file)
        tf_path = os.path.abspath(tf_checkpoint_path)

        print(f"Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.")
        
        if transfo_xl_config_file == "":
            config = TransfoXLConfig()
        else:
            config = TransfoXLConfig.from_json_file(transfo_xl_config_file)
        print(f"Building PyTorch model from configuration: {config}")
        model = TransfoXLLMHeadModel(config)

        model = load_tf_weights_in_transfo_xl(model, config, tf_path)
        
        pytorch_weights_dump_path = os.path.join(pytorch_dump_folder_path, WEIGHTS_NAME)
        pytorch_config_dump_path = os.path.join(pytorch_dump_folder_path, CONFIG_NAME)
        print(f"Save PyTorch model to {os.path.abspath(pytorch_weights_dump_path)}")
        torch.save(model.state_dict(), pytorch_weights_dump_path)
        print(f"Save configuration file to {os.path.abspath(pytorch_config_dump_path)}")
        with open(pytorch_config_dump_path, "w", encoding="utf-8") as f:
            f.write(config.to_json_string())


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        required=True,
        help="Path to the folder to store the PyTorch model or dataset/vocab.",
    )
    parser.add_argument(
        "--tf_checkpoint_path",
        default="",
        type=str,
        help="An optional path to a TensorFlow checkpoint path to be converted.",
    )
    parser.add_argument(
        "--transfo_xl_config_file",
        default="",
        type=str,
        help=(
            "An optional config json file corresponding to the pre-trained BERT model. \n"
            "This specifies the model architecture."
        ),
    )
    parser.add_argument(
        "--transfo_xl_dataset_file",
        default="",
        type=str,
        help="An optional dataset file to be converted in a vocabulary.",
    )
    args = parser.parse_args()
    convert_transfo_xl_checkpoint_to_pytorch(
        args.tf_checkpoint_path,
        args.transfo_xl_config_file,
        args.pytorch_dump_folder_path,
        args.transfo_xl_dataset_file,
    )