codingutf8 2022 google llc 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 t5x checkpoint to pytorch steps install gsutil according to https cloud google comstoragedocsgsutilinstall get a t5x checkpoint at https github comgoogleresearcht5xblobmaindocsmodels mdt511checkpoints example gsutil m cp r gs t5datapretrainedmodelst5xt511small home create or download a corresponding config for the downloaded model e g for t5 v1 1 small you can use https huggingface cogooglet5v11smallblobmainconfig json convert python3 convertt5xcheckpointtopytorch py t5xcheckpointpathhomet511small configfileconfig json pytorchdumppathhomet511smallpt returns the koqv parameters of selfattention does not transpose k paramsfprefixlayersilayernamekeykernel o paramsfprefixlayersilayernameoutkernel q paramsfprefixlayersilayernamequerykernel v paramsfprefixlayersilayernamevaluekernel return k o q v def t5xmlplookupparams i prefix splitmlpwifalse returns the layer norm param of a layer return paramsfprefixlayersilayernamescale def convertt5xtopytorchvariables dict numlayers int numdecoderlayers int isencoderonly bool v1 1 models have a gated gelu with wi0 and wi1 instead of wi shared embeddings encoder block i layer 0 self attention block i layer 1 mlp decoder block i layer 0 self attention block i layer 1 cross attention block i layer 2 mlp lm head only in v1 1 checkpoints in v1 0 embeddings are used instead prepares a state dict for the pytorch model make a state dict with torch tensors statedict collections ordereddictk torch fromnumpyv copy for k v in convertedparams items add what is missing if encoder embedtokens weight not in statedict statedictencoder embedtokens weight statedictshared weight if not isencoderonly if decoder embedtokens weight not in statedict statedictdecoder embedtokens weight statedictshared weight if lmhead weight not in statedict for old 1 0 models printusing shared word embeddings as lmhead statedictlmhead weight statedictshared weight return statedict def loadt5xweightsint5model config t5xcheckpointpath isencoderonly loads the config and model converts the t5x checkpoint and saves a pytorch checkpoint initialise pytorch model config t5config fromjsonfileconfigfile printfbuilding pytorch model from configuration config nonv1 1 checkpoints could also use t5model but this works for all the v1 0 checkpoints will simply have an lm head that is the word embeddings if isencoderonly model t5encodermodelconfig else model t5forconditionalgenerationconfig load weights from tf checkpoint loadt5xweightsint5model config t5xcheckpointpath isencoderonly save pytorchmodel printfsave pytorch model to pytorchdumppath model savepretrainedpytorchdumppath verify that we can load the checkpoint model frompretrainedpytorchdumppath printdone if name main parser argparse argumentparserdescriptionconverts a native t5x checkpoint into a pytorch checkpoint required parameters parser addargument t5xcheckpointpath defaultnone typestr requiredtrue helppath to the t5x checkpoint parser addargument configfile defaultnone typestr requiredtrue helpthe config json file corresponding to the pretrained t5 model nthis specifies the model architecture parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to the output pytorch model parser addargument isencoderonly actionstoretrue helpcheck if the model is encoderdecoder model defaultfalse args parser parseargs convertt5xcheckpointtopytorch args t5xcheckpointpath args configfile args pytorchdumppath args isencoderonly coding utf 8 2022 google llc 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 t5x checkpoint to pytorch steps install gsutil according to https cloud google com storage docs gsutil_install get a t5x checkpoint at https github com google research t5x blob main docs models md t5 11 checkpoints example gsutil m cp r gs t5 data pretrained_models t5x t5_1_1_small home create or download a corresponding config for the downloaded model e g for t5 v1 1 small you can use https huggingface co google t5 v1_1 small blob main config json convert python3 convert_t5x_checkpoint_to_pytorch py t5x_checkpoint_path home t5_1_1_small config_file config json pytorch_dump_path home t5_1_1_small_pt returns the koqv parameters of self attention does not transpose returns the mlp parameters of a layer does not transpose returns the layer norm param of a layer converts the parameters from t5x flax to transformers pytorch v1 1 models have a gated gelu with wi_0 and wi_1 instead of wi shared embeddings encoder block i layer 0 self attention block i layer 1 mlp decoder block i layer 0 self attention block i layer 1 cross attention block i layer 2 mlp lm head only in v1 1 checkpoints in v1 0 embeddings are used instead prepares a state dict for the pytorch model make a state dict with torch tensors add what is missing for old 1 0 models replaces the params in model witht the t5x converted params loads the config and model converts the t5x checkpoint and saves a pytorch checkpoint initialise pytorch model non v1 1 checkpoints could also use t5model but this works for all the v1 0 checkpoints will simply have an lm head that is the word embeddings load weights from tf checkpoint save pytorch model verify that we can load the checkpoint required parameters
import argparse
import collections

import torch
from flax import traverse_util
from t5x import checkpoints

from transformers import T5Config, T5EncoderModel, T5ForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()


def t5x_attention_lookup(params, i, prefix, layer_name="attention"):
    
    k = params[f"{prefix}/layers_{i}/{layer_name}/key/kernel"]
    o = params[f"{prefix}/layers_{i}/{layer_name}/out/kernel"]
    q = params[f"{prefix}/layers_{i}/{layer_name}/query/kernel"]
    v = params[f"{prefix}/layers_{i}/{layer_name}/value/kernel"]
    return k, o, q, v


def t5x_mlp_lookup(params, i, prefix, split_mlp_wi=False):
    
    if split_mlp_wi:
        wi_0 = params[f"{prefix}/layers_{i}/mlp/wi_0/kernel"]
        wi_1 = params[f"{prefix}/layers_{i}/mlp/wi_1/kernel"]
        wi = (wi_0, wi_1)
    else:
        wi = params[f"{prefix}/layers_{i}/mlp/wi/kernel"]

    wo = params[f"{prefix}/layers_{i}/mlp/wo/kernel"]
    return wi, wo


def t5x_layer_norm_lookup(params, i, prefix, layer_name):
    
    return params[f"{prefix}/layers_{i}/{layer_name}/scale"]


def convert_t5x_to_pytorch(variables: dict, *, num_layers: int, num_decoder_layers: int, is_encoder_only: bool):
    
    old = traverse_util.flatten_dict(variables["target"])
    old = {"/".join(k): v for k, v in old.items()}

    
    split_mlp_wi = "encoder/layers_0/mlp/wi_0/kernel" in old
    print("Split MLP:", split_mlp_wi)

    new = collections.OrderedDict()

    
    new["shared.weight"] = old["token_embedder/embedding"]

    
    for i in range(num_layers):
        
        layer_norm = t5x_layer_norm_lookup(old, i, "encoder", "pre_attention_layer_norm")
        k, o, q, v = t5x_attention_lookup(old, i, "encoder", "attention")
        new[f"encoder.block.{i}.layer.0.layer_norm.weight"] = layer_norm
        new[f"encoder.block.{i}.layer.0.SelfAttention.k.weight"] = k.T
        new[f"encoder.block.{i}.layer.0.SelfAttention.o.weight"] = o.T
        new[f"encoder.block.{i}.layer.0.SelfAttention.q.weight"] = q.T
        new[f"encoder.block.{i}.layer.0.SelfAttention.v.weight"] = v.T

        
        layer_norm = t5x_layer_norm_lookup(old, i, "encoder", "pre_mlp_layer_norm")
        wi, wo = t5x_mlp_lookup(old, i, "encoder", split_mlp_wi)
        new[f"encoder.block.{i}.layer.1.layer_norm.weight"] = layer_norm
        if split_mlp_wi:
            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi_0.weight"] = wi[0].T
            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi_1.weight"] = wi[1].T
        else:
            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi.weight"] = wi.T
        new[f"encoder.block.{i}.layer.1.DenseReluDense.wo.weight"] = wo.T

    new["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = old[
        "encoder/relpos_bias/rel_embedding"
    ].T
    new["encoder.final_layer_norm.weight"] = old["encoder/encoder_norm/scale"]

    if not is_encoder_only:
        
        for i in range(num_decoder_layers):
            
            layer_norm = t5x_layer_norm_lookup(old, i, "decoder", "pre_self_attention_layer_norm")
            k, o, q, v = t5x_attention_lookup(old, i, "decoder", "self_attention")
            new[f"decoder.block.{i}.layer.0.layer_norm.weight"] = layer_norm
            new[f"decoder.block.{i}.layer.0.SelfAttention.k.weight"] = k.T
            new[f"decoder.block.{i}.layer.0.SelfAttention.o.weight"] = o.T
            new[f"decoder.block.{i}.layer.0.SelfAttention.q.weight"] = q.T
            new[f"decoder.block.{i}.layer.0.SelfAttention.v.weight"] = v.T

            
            layer_norm = t5x_layer_norm_lookup(old, i, "decoder", "pre_cross_attention_layer_norm")
            k, o, q, v = t5x_attention_lookup(old, i, "decoder", "encoder_decoder_attention")
            new[f"decoder.block.{i}.layer.1.layer_norm.weight"] = layer_norm
            new[f"decoder.block.{i}.layer.1.EncDecAttention.k.weight"] = k.T
            new[f"decoder.block.{i}.layer.1.EncDecAttention.o.weight"] = o.T
            new[f"decoder.block.{i}.layer.1.EncDecAttention.q.weight"] = q.T
            new[f"decoder.block.{i}.layer.1.EncDecAttention.v.weight"] = v.T

            
            layer_norm = t5x_layer_norm_lookup(old, i, "decoder", "pre_mlp_layer_norm")
            wi, wo = t5x_mlp_lookup(old, i, "decoder", split_mlp_wi)
            new[f"decoder.block.{i}.layer.2.layer_norm.weight"] = layer_norm
            if split_mlp_wi:
                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi_0.weight"] = wi[0].T
                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi_1.weight"] = wi[1].T
            else:
                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi.weight"] = wi.T
            new[f"decoder.block.{i}.layer.2.DenseReluDense.wo.weight"] = wo.T

        new["decoder.final_layer_norm.weight"] = old["decoder/decoder_norm/scale"]
        new["decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = old[
            "decoder/relpos_bias/rel_embedding"
        ].T

        
        if "decoder/logits_dense/kernel" in old:
            new["lm_head.weight"] = old["decoder/logits_dense/kernel"].T

    return new


def make_state_dict(converted_params, is_encoder_only: bool):
    
    
    state_dict = collections.OrderedDict([(k, torch.from_numpy(v.copy())) for (k, v) in converted_params.items()])

    
    if "encoder.embed_tokens.weight" not in state_dict:
        state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"]

    if not is_encoder_only:
        if "decoder.embed_tokens.weight" not in state_dict:
            state_dict["decoder.embed_tokens.weight"] = state_dict["shared.weight"]

        if "lm_head.weight" not in state_dict:  
            print("Using shared word embeddings as lm_head.")
            state_dict["lm_head.weight"] = state_dict["shared.weight"]

    return state_dict


def load_t5x_weights_in_t5(model, config, t5x_checkpoint_path, is_encoder_only):
    
    variables = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path)
    converted = convert_t5x_to_pytorch(
        variables,
        num_layers=config.num_layers,
        num_decoder_layers=config.num_decoder_layers,
        is_encoder_only=is_encoder_only,
    )
    state_dict = make_state_dict(converted, is_encoder_only)
    model.load_state_dict(state_dict, strict=True)


def convert_t5x_checkpoint_to_pytorch(
    t5x_checkpoint_path, config_file, pytorch_dump_path, is_encoder_only: bool = False
):
    
    
    config = T5Config.from_json_file(config_file)
    print(f"Building PyTorch model from configuration: {config}")
    
    
    if is_encoder_only:
        model = T5EncoderModel(config)
    else:
        model = T5ForConditionalGeneration(config)

    
    load_t5x_weights_in_t5(model, config, t5x_checkpoint_path, is_encoder_only)

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

    
    model.from_pretrained(pytorch_dump_path)
    print("Done")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Converts a native T5X checkpoint into a PyTorch checkpoint.")
    
    parser.add_argument(
        "--t5x_checkpoint_path", default=None, type=str, required=True, help="Path to the T5X checkpoint."
    )
    parser.add_argument(
        "--config_file",
        default=None,
        type=str,
        required=True,
        help="The config json file corresponding to the pre-trained T5 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(
        "--is_encoder_only", action="store_true", help="Check if the model is encoder-decoder model", default=False
    )
    args = parser.parse_args()
    convert_t5x_checkpoint_to_pytorch(
        args.t5x_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
    )