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 canine checkpoint import argparse from transformers import canineconfig caninemodel caninetokenizer loadtfweightsincanine from transformers utils import logging logging setverbosityinfo def converttfcheckpointtopytorchtfcheckpointpath pytorchdumppath initialize pytorch model config canineconfig model caninemodelconfig model eval printfbuilding pytorch model from configuration config load weights from tf checkpoint loadtfweightsincaninemodel config tfcheckpointpath save pytorchmodel weights and configuration printfsave pytorch model to pytorchdumppath model savepretrainedpytorchdumppath save tokenizer files tokenizer caninetokenizer printfsave tokenizer files to pytorchdumppath tokenizer savepretrainedpytorchdumppath if name main parser argparse argumentparser required parameters parser addargument tfcheckpointpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint should end with model ckpt parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to a folder where the pytorch model will be placed args parser parseargs converttfcheckpointtopytorchargs tfcheckpointpath args pytorchdumppath 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 canine checkpoint initialize pytorch model load weights from tf checkpoint save pytorch model weights and configuration save tokenizer files required parameters
import argparse

from transformers import CanineConfig, CanineModel, CanineTokenizer, load_tf_weights_in_canine
from transformers.utils import logging


logging.set_verbosity_info()


def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, pytorch_dump_path):
    
    config = CanineConfig()
    model = CanineModel(config)
    model.eval()

    print(f"Building PyTorch model from configuration: {config}")

    
    load_tf_weights_in_canine(model, config, tf_checkpoint_path)

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

    
    tokenizer = CanineTokenizer()
    print(f"Save tokenizer files to {pytorch_dump_path}")
    tokenizer.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. Should end with model.ckpt",
    )
    parser.add_argument(
        "--pytorch_dump_path",
        default=None,
        type=str,
        required=True,
        help="Path to a folder where the PyTorch model will be placed.",
    )
    args = parser.parse_args()
    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.pytorch_dump_path)