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 electra checkpoint import argparse import torch from transformers import electraconfig electraformaskedlm electraforpretraining loadtfweightsinelectra from transformers utils import logging logging setverbosityinfo def converttfcheckpointtopytorchtfcheckpointpath configfile pytorchdumppath discriminatororgenerator initialise pytorch model config electraconfig fromjsonfileconfigfile printfbuilding pytorch model from configuration config if discriminatororgenerator discriminator model electraforpretrainingconfig elif discriminatororgenerator generator model electraformaskedlmconfig else raise valueerrorthe discriminatororgenerator argument should be either discriminator or generator load weights from tf checkpoint loadtfweightsinelectra model config tfcheckpointpath discriminatororgeneratordiscriminatororgenerator save pytorchmodel printfsave pytorch model to pytorchdumppath torch savemodel statedict pytorchdumppath 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 helpthe config json file corresponding to the pretrained model nthis specifies the model architecture parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to the output pytorch model parser addargument discriminatororgenerator defaultnone typestr requiredtrue help whether to export the generator or the discriminator should be a string either discriminator or generator args parser parseargs converttfcheckpointtopytorch args tfcheckpointpath args configfile args pytorchdumppath args discriminatororgenerator 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 electra checkpoint initialise pytorch model load weights from tf checkpoint save pytorch model required parameters
import argparse

import torch

from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
from transformers.utils import logging


logging.set_verbosity_info()


def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path, discriminator_or_generator):
    
    config = ElectraConfig.from_json_file(config_file)
    print(f"Building PyTorch model from configuration: {config}")

    if discriminator_or_generator == "discriminator":
        model = ElectraForPreTraining(config)
    elif discriminator_or_generator == "generator":
        model = ElectraForMaskedLM(config)
    else:
        raise ValueError("The discriminator_or_generator argument should be either 'discriminator' or 'generator'")

    
    load_tf_weights_in_electra(
        model, config, tf_checkpoint_path, discriminator_or_generator=discriminator_or_generator
    )

    
    print(f"Save PyTorch model to {pytorch_dump_path}")
    torch.save(model.state_dict(), 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 model. \nThis specifies the model architecture.",
    )
    parser.add_argument(
        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--discriminator_or_generator",
        default=None,
        type=str,
        required=True,
        help=(
            "Whether to export the generator or the discriminator. Should be a string, either 'discriminator' or "
            "'generator'."
        ),
    )
    args = parser.parse_args()
    convert_tf_checkpoint_to_pytorch(
        args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
    )