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 bert checkpoint import argparse import os import torch from transformers import xlnetconfig xlnetforquestionanswering xlnetforsequenceclassification xlnetlmheadmodel loadtfweightsinxlnet from transformers utils import configname weightsname logging gluetasksnumlabels cola 2 mnli 3 mrpc 2 sst2 2 stsb 1 qqp 2 qnli 2 rte 2 wnli 2 logging setverbosityinfo def convertxlnetcheckpointtopytorch tfcheckpointpath bertconfigfile pytorchdumpfolderpath finetuningtasknone initialise pytorch model config xlnetconfig fromjsonfilebertconfigfile finetuningtask finetuningtask lower if finetuningtask is not none else if finetuningtask in gluetasksnumlabels printfbuilding pytorch xlnetforsequenceclassification model from configuration config config finetuningtask finetuningtask config numlabels gluetasksnumlabelsfinetuningtask model xlnetforsequenceclassificationconfig elif squad in finetuningtask config finetuningtask finetuningtask model xlnetforquestionansweringconfig else model xlnetlmheadmodelconfig load weights from tf checkpoint loadtfweightsinxlnetmodel config tfcheckpointpath 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 required parameters parser addargument tfcheckpointpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument xlnetconfigfile defaultnone typestr requiredtrue help the config json file corresponding to the pretrained xlnet model n this specifies the model architecture parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the folder to store the pytorch model or datasetvocab parser addargument finetuningtask defaultnone typestr helpname of a task on which the xlnet tensorflow model was finetuned args parser parseargs printargs convertxlnetcheckpointtopytorch args tfcheckpointpath args xlnetconfigfile args pytorchdumpfolderpath args finetuningtask 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 bert checkpoint initialise pytorch model load weights from tf checkpoint save pytorch model required parameters
import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


GLUE_TASKS_NUM_LABELS = {
    "cola": 2,
    "mnli": 3,
    "mrpc": 2,
    "sst-2": 2,
    "sts-b": 1,
    "qqp": 2,
    "qnli": 2,
    "rte": 2,
    "wnli": 2,
}


logging.set_verbosity_info()


def convert_xlnet_checkpoint_to_pytorch(
    tf_checkpoint_path, bert_config_file, pytorch_dump_folder_path, finetuning_task=None
):
    
    config = XLNetConfig.from_json_file(bert_config_file)

    finetuning_task = finetuning_task.lower() if finetuning_task is not None else ""
    if finetuning_task in GLUE_TASKS_NUM_LABELS:
        print(f"Building PyTorch XLNetForSequenceClassification model from configuration: {config}")
        config.finetuning_task = finetuning_task
        config.num_labels = GLUE_TASKS_NUM_LABELS[finetuning_task]
        model = XLNetForSequenceClassification(config)
    elif "squad" in finetuning_task:
        config.finetuning_task = finetuning_task
        model = XLNetForQuestionAnswering(config)
    else:
        model = XLNetLMHeadModel(config)

    
    load_tf_weights_in_xlnet(model, config, tf_checkpoint_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(
        "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
    )
    parser.add_argument(
        "--xlnet_config_file",
        default=None,
        type=str,
        required=True,
        help=(
            "The config json file corresponding to the pre-trained XLNet model. \n"
            "This specifies the model architecture."
        ),
    )
    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(
        "--finetuning_task",
        default=None,
        type=str,
        help="Name of a task on which the XLNet TensorFlow model was fine-tuned",
    )
    args = parser.parse_args()
    print(args)

    convert_xlnet_checkpoint_to_pytorch(
        args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
    )