codingutf8 2021 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 image gpt checkpoints import argparse import torch from transformers import imagegptconfig imagegptforcausallm loadtfweightsinimagegpt from transformers utils import configname weightsname logging logging setverbosityinfo def convertimagegptcheckpointtopytorchimagegptcheckpointpath modelsize pytorchdumpfolderpath construct configuration depending on size models small 512 8 24 medium 1024 8 36 large 1536 16 48 nembd nhead nlayer modelsmodelsize set model hyperparameters config imagegptconfignembdnembd nlayernlayer nheadnhead model imagegptforcausallmconfig load weights from numpy loadtfweightsinimagegptmodel config imagegptcheckpointpath 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 imagegptcheckpointpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument modelsize defaultnone typestr requiredtrue helpsize of the model can be either small medium or large parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs convertimagegptcheckpointtopytorch args imagegptcheckpointpath args modelsize args pytorchdumpfolderpath coding utf 8 2021 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 image gpt checkpoints construct configuration depending on size set model hyperparameters load weights from numpy save pytorch model required parameters
import argparse

import torch

from transformers import ImageGPTConfig, ImageGPTForCausalLM, load_tf_weights_in_imagegpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()


def convert_imagegpt_checkpoint_to_pytorch(imagegpt_checkpoint_path, model_size, pytorch_dump_folder_path):
    
    MODELS = {"small": (512, 8, 24), "medium": (1024, 8, 36), "large": (1536, 16, 48)}
    n_embd, n_head, n_layer = MODELS[model_size]  
    config = ImageGPTConfig(n_embd=n_embd, n_layer=n_layer, n_head=n_head)
    model = ImageGPTForCausalLM(config)

    
    load_tf_weights_in_imagegpt(model, config, imagegpt_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(
        "--imagegpt_checkpoint_path",
        default=None,
        type=str,
        required=True,
        help="Path to the TensorFlow checkpoint path.",
    )
    parser.add_argument(
        "--model_size",
        default=None,
        type=str,
        required=True,
        help="Size of the model (can be either 'small', 'medium' or 'large').",
    )
    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_imagegpt_checkpoint_to_pytorch(
        args.imagegpt_checkpoint_path, args.model_size, args.pytorch_dump_folder_path
    )