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 gpt2config gpt2model loadtfweightsingpt2 from transformers utils import configname weightsname logging logging setverbosityinfo def convertgpt2checkpointtopytorchgpt2checkpointpath gpt2configfile pytorchdumpfolderpath construct model if gpt2configfile config gpt2config else config gpt2config fromjsonfilegpt2configfile model gpt2modelconfig load weights from numpy loadtfweightsingpt2model config gpt2checkpointpath 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 gpt2checkpointpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model parser addargument gpt2configfile default typestr help an optional config json file corresponding to the pretrained openai model n this specifies the model architecture args parser parseargs convertgpt2checkpointtopytorchargs gpt2checkpointpath args gpt2configfile 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 GPT2Config, GPT2Model, load_tf_weights_in_gpt2
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()


def convert_gpt2_checkpoint_to_pytorch(gpt2_checkpoint_path, gpt2_config_file, pytorch_dump_folder_path):
    
    if gpt2_config_file == "":
        config = GPT2Config()
    else:
        config = GPT2Config.from_json_file(gpt2_config_file)
    model = GPT2Model(config)

    
    load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_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(
        "--gpt2_checkpoint_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(
        "--gpt2_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_gpt2_checkpoint_to_pytorch(args.gpt2_checkpoint_path, args.gpt2_config_file, args.pytorch_dump_folder_path)