2020 the huggingface 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 this script can be used to convert a headless tf2 x bert model to pytorch as published on the official now deprecated github https github comtensorflowmodelstreev2 3 0officialnlpbert tf2 x uses different variable names from the original bert tf 1 4 implementation the script remaps the tf2 x bert weight names to the original names so the model can be imported with huggingfacetransformer you may adapt this script to include classificationmlmnspetc heads note this script is only working with an older version of the tensorflow models repository v2 3 0 models trained with never versions are not compatible with this script load weights from tf model logger infofloading tf weight name with shape shape ignore initial model figure out how many levels deep the name is read data sanity check convert layers variable names end with attributesvariablevalue embedding layers layernum 0 wordembeddings layernum 1 positionembeddings layernum 2 tokentypeembeddings embedding layernorm encoder layers pooler layer selfattention layer output attention norm output attention dense output dense output dense attention key attention query attention value attention intermediate dense output layer norm weights biases for certain layers reshape is necessary instantiate model load weights from checkpoint save pytorchmodel 2020 the huggingface 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 this script can be used to convert a head less tf2 x bert model to pytorch as published on the official now deprecated github https github com tensorflow models tree v2 3 0 official nlp bert tf2 x uses different variable names from the original bert tf 1 4 implementation the script re maps the tf2 x bert weight names to the original names so the model can be imported with huggingface transformer you may adapt this script to include classification mlm nsp etc heads note this script is only working with an older version of the tensorflow models repository v2 3 0 models trained with never versions are not compatible with this script load weights from tf model logger info f loading tf weight name with shape shape ignore initial model figure out how many levels deep the name is read data sanity check convert layers variable names end with attributes variable_value embedding layers layer_num 0 word_embeddings layer_num 1 position_embeddings layer_num 2 token_type_embeddings embedding layernorm encoder layers pooler layer self attention layer output attention norm output attention dense output dense output dense attention key attention query attention value attention intermediate dense output layer norm weights biases for certain layers reshape is necessary instantiate model load weights from checkpoint save pytorch model
import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
logger = logging.get_logger(__name__)


def load_tf2_weights_in_bert(model, tf_checkpoint_path, config):
    tf_path = os.path.abspath(tf_checkpoint_path)
    logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
    
    init_vars = tf.train.list_variables(tf_path)
    names = []
    arrays = []
    layer_depth = []
    for full_name, shape in init_vars:
        
        name = full_name.split("/")
        if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
            logger.info(f"Skipping non-model layer {full_name}")
            continue
        if "optimizer" in full_name:
            logger.info(f"Skipping optimization layer {full_name}")
            continue
        if name[0] == "model":
            
            name = name[1:]
        
        depth = 0
        for _name in name:
            if _name.startswith("layer_with_weights"):
                depth += 1
            else:
                break
        layer_depth.append(depth)
        
        array = tf.train.load_variable(tf_path, full_name)
        names.append("/".join(name))
        arrays.append(array)
    logger.info(f"Read a total of {len(arrays):,} layers")

    
    if len(set(layer_depth)) != 1:
        raise ValueError(f"Found layer names with different depths (layer depth {list(set(layer_depth))})")
    layer_depth = list(set(layer_depth))[0]
    if layer_depth != 1:
        raise ValueError(
            "The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP"
            " heads."
        )

    
    logger.info("Converting weights...")
    for full_name, array in zip(names, arrays):
        name = full_name.split("/")
        pointer = model
        trace = []
        for i, m_name in enumerate(name):
            if m_name == ".ATTRIBUTES":
                
                break
            if m_name.startswith("layer_with_weights"):
                layer_num = int(m_name.split("-")[-1])
                if layer_num <= 2:
                    
                    
                    
                    
                    continue
                elif layer_num == 3:
                    
                    trace.extend(["embeddings", "LayerNorm"])
                    pointer = getattr(pointer, "embeddings")
                    pointer = getattr(pointer, "LayerNorm")
                elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
                    
                    trace.extend(["encoder", "layer", str(layer_num - 4)])
                    pointer = getattr(pointer, "encoder")
                    pointer = getattr(pointer, "layer")
                    pointer = pointer[layer_num - 4]
                elif layer_num == config.num_hidden_layers + 4:
                    
                    trace.extend(["pooler", "dense"])
                    pointer = getattr(pointer, "pooler")
                    pointer = getattr(pointer, "dense")
            elif m_name == "embeddings":
                trace.append("embeddings")
                pointer = getattr(pointer, "embeddings")
                if layer_num == 0:
                    trace.append("word_embeddings")
                    pointer = getattr(pointer, "word_embeddings")
                elif layer_num == 1:
                    trace.append("position_embeddings")
                    pointer = getattr(pointer, "position_embeddings")
                elif layer_num == 2:
                    trace.append("token_type_embeddings")
                    pointer = getattr(pointer, "token_type_embeddings")
                else:
                    raise ValueError(f"Unknown embedding layer with name {full_name}")
                trace.append("weight")
                pointer = getattr(pointer, "weight")
            elif m_name == "_attention_layer":
                
                trace.extend(["attention", "self"])
                pointer = getattr(pointer, "attention")
                pointer = getattr(pointer, "self")
            elif m_name == "_attention_layer_norm":
                
                trace.extend(["attention", "output", "LayerNorm"])
                pointer = getattr(pointer, "attention")
                pointer = getattr(pointer, "output")
                pointer = getattr(pointer, "LayerNorm")
            elif m_name == "_attention_output_dense":
                
                trace.extend(["attention", "output", "dense"])
                pointer = getattr(pointer, "attention")
                pointer = getattr(pointer, "output")
                pointer = getattr(pointer, "dense")
            elif m_name == "_output_dense":
                
                trace.extend(["output", "dense"])
                pointer = getattr(pointer, "output")
                pointer = getattr(pointer, "dense")
            elif m_name == "_output_layer_norm":
                
                trace.extend(["output", "LayerNorm"])
                pointer = getattr(pointer, "output")
                pointer = getattr(pointer, "LayerNorm")
            elif m_name == "_key_dense":
                
                trace.append("key")
                pointer = getattr(pointer, "key")
            elif m_name == "_query_dense":
                
                trace.append("query")
                pointer = getattr(pointer, "query")
            elif m_name == "_value_dense":
                
                trace.append("value")
                pointer = getattr(pointer, "value")
            elif m_name == "_intermediate_dense":
                
                trace.extend(["intermediate", "dense"])
                pointer = getattr(pointer, "intermediate")
                pointer = getattr(pointer, "dense")
            elif m_name == "_output_layer_norm":
                
                trace.append("output")
                pointer = getattr(pointer, "output")
            
            elif m_name in ["bias", "beta"]:
                trace.append("bias")
                pointer = getattr(pointer, "bias")
            elif m_name in ["kernel", "gamma"]:
                trace.append("weight")
                pointer = getattr(pointer, "weight")
            else:
                logger.warning(f"Ignored {m_name}")
        
        trace = ".".join(trace)
        if re.match(r"(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)", trace) or re.match(
            r"(\S+)\.attention\.output\.dense\.weight", trace
        ):
            array = array.reshape(pointer.data.shape)
        if "kernel" in full_name:
            array = array.transpose()
        if pointer.shape == array.shape:
            pointer.data = torch.from_numpy(array)
        else:
            raise ValueError(
                f"Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:"
                f" {array.shape}"
            )
        logger.info(f"Successfully set variable {full_name} to PyTorch layer {trace}")
    return model


def convert_tf2_checkpoint_to_pytorch(tf_checkpoint_path, config_path, pytorch_dump_path):
    
    logger.info(f"Loading model based on config from {config_path}...")
    config = BertConfig.from_json_file(config_path)
    model = BertModel(config)

    
    logger.info(f"Loading weights from checkpoint {tf_checkpoint_path}...")
    load_tf2_weights_in_bert(model, tf_checkpoint_path, config)

    
    logger.info(f"Saving 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", type=str, required=True, help="Path to the TensorFlow 2.x checkpoint path."
    )
    parser.add_argument(
        "--bert_config_file",
        type=str,
        required=True,
        help="The config json file corresponding to the BERT model. This specifies the model architecture.",
    )
    parser.add_argument(
        "--pytorch_dump_path",
        type=str,
        required=True,
        help="Path to the output PyTorch model (must include filename).",
    )
    args = parser.parse_args()
    convert_tf2_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)