codingutf8 2021 the eleuther ai and 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 gpt neo checkpoint import argparse import json from transformers import gptneoconfig gptneoforcausallm loadtfweightsingptneo from transformers utils import logging logging setverbosityinfo def converttfcheckpointtopytorchtfcheckpointpath configfile pytorchdumppath initialise pytorch model configjson json loadopenconfigfile r config gptneoconfig hiddensizeconfigjsonnembd numlayersconfigjsonnlayer numheadsconfigjsonnhead attentiontypesconfigjsonattentiontypes maxpositionembeddingsconfigjsonnpositions residdropoutconfigjsonresdropout embeddropoutconfigjsonembeddropout attentiondropoutconfigjsonattndropout printfbuilding pytorch model from configuration config model gptneoforcausallmconfig load weights from tf checkpoint loadtfweightsingptneomodel config tfcheckpointpath save pytorchmodel printfsave pytorch model to pytorchdumppath model savepretrainedpytorchdumppath if name main parser argparse argumentparser required parameters parser addargument tfcheckpointpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument configfile defaultnone typestr requiredtrue help the config json file corresponding to the pretrained meshtf model n this specifies the model architecture parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs converttfcheckpointtopytorchargs tfcheckpointpath args configfile args pytorchdumppath coding utf 8 2021 the eleuther ai and 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 gpt neo checkpoint initialise pytorch model load weights from tf checkpoint save pytorch model required parameters
import argparse
import json

from transformers import GPTNeoConfig, GPTNeoForCausalLM, load_tf_weights_in_gpt_neo
from transformers.utils import logging


logging.set_verbosity_info()


def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path):
    
    config_json = json.load(open(config_file, "r"))
    config = GPTNeoConfig(
        hidden_size=config_json["n_embd"],
        num_layers=config_json["n_layer"],
        num_heads=config_json["n_head"],
        attention_types=config_json["attention_types"],
        max_position_embeddings=config_json["n_positions"],
        resid_dropout=config_json["res_dropout"],
        embed_dropout=config_json["embed_dropout"],
        attention_dropout=config_json["attn_dropout"],
    )
    print(f"Building PyTorch model from configuration: {config}")
    model = GPTNeoForCausalLM(config)

    
    load_tf_weights_in_gpt_neo(model, config, tf_checkpoint_path)

    
    print(f"Save PyTorch model to {pytorch_dump_path}")
    model.save_pretrained(pytorch_dump_path)


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(
        "--config_file",
        default=None,
        type=str,
        required=True,
        help=(
            "The config json file corresponding to the pre-trained mesh-tf model. \n"
            "This specifies the model architecture."
        ),
    )
    parser.add_argument(
        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    args = parser.parse_args()
    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)