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 torch from transformers import openaigptconfig openaigptmodel loadtfweightsinopenaigpt from transformers utils import configname weightsname logging logging setverbosityinfo def convertopenaicheckpointtopytorchopenaicheckpointfolderpath openaiconfigfile pytorchdumpfolderpath construct model if openaiconfigfile config openaigptconfig else config openaigptconfig fromjsonfileopenaiconfigfile model openaigptmodelconfig load weights from numpy loadtfweightsinopenaigptmodel config openaicheckpointfolderpath save pytorchmodel pytorchweightsdumppath pytorchdumpfolderpath weightsname pytorchconfigdumppath pytorchdumpfolderpath configname printfsave pytorch model to pytorchweightsdumppath torch savemodel statedict pytorchweightsdumppath printfsave configuration file to pytorchconfigdumppath with openpytorchconfigdumppath w encodingutf8 as f f writeconfig tojsonstring if name main parser argparse argumentparser required parameters parser addargument openaicheckpointfolderpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model parser addargument openaiconfigfile default typestr help an optional config json file corresponding to the pretrained openai model n this specifies the model architecture args parser parseargs convertopenaicheckpointtopytorch args openaicheckpointfolderpath args openaiconfigfile 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 construct model load weights from numpy save pytorch model required parameters
import argparse

import torch

from transformers import OpenAIGPTConfig, OpenAIGPTModel, load_tf_weights_in_openai_gpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()


def convert_openai_checkpoint_to_pytorch(openai_checkpoint_folder_path, openai_config_file, pytorch_dump_folder_path):
    
    if openai_config_file == "":
        config = OpenAIGPTConfig()
    else:
        config = OpenAIGPTConfig.from_json_file(openai_config_file)
    model = OpenAIGPTModel(config)

    
    load_tf_weights_in_openai_gpt(model, config, openai_checkpoint_folder_path)

    
    pytorch_weights_dump_path = pytorch_dump_folder_path + "/" + WEIGHTS_NAME
    pytorch_config_dump_path = pytorch_dump_folder_path + "/" + CONFIG_NAME
    print(f"Save PyTorch model to {pytorch_weights_dump_path}")
    torch.save(model.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(config.to_json_string())


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--openai_checkpoint_folder_path",
        default=None,
        type=str,
        required=True,
        help="Path to the TensorFlow checkpoint path.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--openai_config_file",
        default="",
        type=str,
        help=(
            "An optional config json file corresponding to the pre-trained OpenAI model. \n"
            "This specifies the model architecture."
        ),
    )
    args = parser.parse_args()
    convert_openai_checkpoint_to_pytorch(
        args.openai_checkpoint_folder_path, args.openai_config_file, args.pytorch_dump_folder_path
    )