codingutf8 2020 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 seq2seq tf hub checkpoint import argparse from import bertconfig bertgenerationconfig bertgenerationdecoder bertgenerationencoder loadtfweightsinbertgeneration logging logging setverbosityinfo def converttfcheckpointtopytorchtfhubpath pytorchdumppath isencodernameddecoder vocabsize isencoder initialise pytorch model bertconfig bertconfig frompretrained bertlargecased vocabsizevocabsize maxpositionembeddings512 isdecodertrue addcrossattentiontrue bertconfigdict bertconfig todict del bertconfigdicttypevocabsize config bertgenerationconfigbertconfigdict if isencoder model bertgenerationencoderconfig else model bertgenerationdecoderconfig printfbuilding pytorch model from configuration config load weights from tf checkpoint loadtfweightsinbertgeneration model tfhubpath modelclassbert isencodernameddecoderisencodernameddecoder isencoderisencoder save pytorchmodel printfsave pytorch model and config to pytorchdumppath model savepretrainedpytorchdumppath if name main parser argparse argumentparser required parameters parser addargument tfhubpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to the output pytorch model parser addargument isencodernameddecoder actionstoretrue helpif decoder has to be renamed to encoder in pytorch model parser addargumentisencoder actionstoretrue helpif model is an encoder parser addargumentvocabsize default50358 typeint helpvocab size of model args parser parseargs converttfcheckpointtopytorch args tfhubpath args pytorchdumppath args isencodernameddecoder args vocabsize isencoderargs isencoder coding utf 8 2020 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 seq2seq tf hub checkpoint initialise pytorch model load weights from tf checkpoint save pytorch model required parameters
import argparse

from . import (
    BertConfig,
    BertGenerationConfig,
    BertGenerationDecoder,
    BertGenerationEncoder,
    load_tf_weights_in_bert_generation,
    logging,
)


logging.set_verbosity_info()


def convert_tf_checkpoint_to_pytorch(tf_hub_path, pytorch_dump_path, is_encoder_named_decoder, vocab_size, is_encoder):
    
    bert_config = BertConfig.from_pretrained(
        "bert-large-cased",
        vocab_size=vocab_size,
        max_position_embeddings=512,
        is_decoder=True,
        add_cross_attention=True,
    )
    bert_config_dict = bert_config.to_dict()
    del bert_config_dict["type_vocab_size"]
    config = BertGenerationConfig(**bert_config_dict)
    if is_encoder:
        model = BertGenerationEncoder(config)
    else:
        model = BertGenerationDecoder(config)
    print(f"Building PyTorch model from configuration: {config}")

    
    load_tf_weights_in_bert_generation(
        model,
        tf_hub_path,
        model_class="bert",
        is_encoder_named_decoder=is_encoder_named_decoder,
        is_encoder=is_encoder,
    )

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


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--tf_hub_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
    )
    parser.add_argument(
        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--is_encoder_named_decoder",
        action="store_true",
        help="If decoder has to be renamed to encoder in PyTorch model.",
    )
    parser.add_argument("--is_encoder", action="store_true", help="If model is an encoder.")
    parser.add_argument("--vocab_size", default=50358, type=int, help="Vocab size of model")
    args = parser.parse_args()
    convert_tf_checkpoint_to_pytorch(
        args.tf_hub_path,
        args.pytorch_dump_path,
        args.is_encoder_named_decoder,
        args.vocab_size,
        is_encoder=args.is_encoder,
    )