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 reformer checkpoint import argparse import pickle import numpy as np import torch from torch import nn from transformers import reformerconfig reformermodelwithlmhead from transformers utils import logging logging setverbosityinfo def setparamtorchlayer weight biasnone set parameter of one layer assert torchlayer weight shape weight shape ftorchlayer layer weight does not match torchlayer weight nn parameterweight if bias is not none assert torchlayer bias shape bias shape ftorchlayer layer bias does not match torchlayer bias nn parameterbias def setlayerweightsintorchlshweights torchlayer hiddensize set torch weights for 1to1 comparison npquerykey np asarrayweights0 npvalue np asarrayweights1 npdense np asarrayweights2 setparam torchlayer selfattention querykey torch tensornpquerykey transpose1 2 contiguous view1 hiddensize setparam torchlayer selfattention value torch tensornpvalue transpose1 2 contiguous view1 hiddensize setparam torchlayer output dense torch tensornpdense view1 hiddensize contiguous transpose0 1 def setlayerweightsintorchlocalweights torchlayer hiddensize set torch weights for 1to1 comparison npquery np asarrayweights0 npkey np asarrayweights1 npvalue np asarrayweights2 npdense np asarrayweights3 setparam torchlayer selfattention query torch tensornpquery transpose1 2 contiguous view1 hiddensize setparam torchlayer selfattention key torch tensornpkey transpose1 2 contiguous view1 hiddensize setparam torchlayer selfattention value torch tensornpvalue transpose1 2 contiguous view1 hiddensize setparam torchlayer output dense torch tensornpdense view1 hiddensize contiguous transpose0 1 def setblockweightsintorchweights torchblock hiddensize layernorm 1 layernorm1 weights000 layernorm1weight np asarraylayernorm10 layernorm1bias np asarraylayernorm11 setparam torchblock attention layernorm torch tensorlayernorm1weight torch tensorlayernorm1bias lsh weights output attnweights weights01 if lenattnweights 4 setlayerweightsintorchlshattnweights torchblock attention hiddensize else setlayerweightsintorchlocalattnweights torchblock attention hiddensize intermediate weighs intermediateweights weights2012 chunked feed forward if lenintermediateweights 4 intermediateweights intermediateweights2 layernorm 2 layernorm2weight np asarrayintermediateweights00 layernorm2bias np asarrayintermediateweights01 setparam torchblock feedforward layernorm torch tensorlayernorm2weight torch tensorlayernorm2bias intermediate dense interdenseweight np asarrayintermediateweights10 interdensebias np asarrayintermediateweights11 setparam torchblock feedforward dense dense torch tensorinterdenseweight transpose0 1 contiguous torch tensorinterdensebias intermediate out outdenseweight np asarrayintermediateweights40 outdensebias np asarrayintermediateweights41 setparam torchblock feedforward output dense torch tensoroutdenseweight transpose0 1 contiguous torch tensoroutdensebias def setmodelweightsintorchweights torchmodel hiddensize reformer model torchmodelreformer torchmodel reformer word embeds wordembeddings np asarrayweights1 setparam torchmodelreformer embeddings wordembeddings torch tensorwordembeddings if isinstanceweights3 tuple positionembeddings torchmodelreformer embeddings positionembeddings for embidx in rangelenpositionembeddings weights embweights np asarrayweights3embidx0 assert positionembeddings weightsembidx shape embweights shape fpositionembeddingsembidx emb does not match positionembeddings weightsembidx nn parametertorch tensorembweights traxlayerweights weights5 assert lentorchmodelreformer encoder layers 4 len traxlayerweights hf and trax model do not have the same number of layers for layeridx layer in enumeratetorchmodelreformer encoder layers blockweights traxlayerweights4 layeridx 4 layeridx 1 setblockweightsintorchblockweights layer hiddensize output layer norm layernormoutweight np asarrayweights70 layernormoutbias np asarrayweights71 setparam torchmodelreformer encoder layernorm torch tensorlayernormoutweight torch tensorlayernormoutbias output embeddings outputembedweights np asarrayweights90 outputembedbias np asarrayweights91 setparam torchmodel lmhead decoder torch tensoroutputembedweights transpose0 1 contiguous torch tensoroutputembedbias def converttraxcheckpointtopytorchtraxmodelpklpath configfile pytorchdumppath initialise pytorch model config reformerconfig fromjsonfileconfigfile printfbuilding pytorch model from configuration config model reformermodelwithlmheadconfig with opentraxmodelpklpath rb as f modelweights pickle loadfweights setmodelweightsintorchmodelweights model config hiddensize save pytorchmodel printfsave pytorch model to pytorchdumppath torch savemodel statedict pytorchdumppath if name main parser argparse argumentparser required parameters parser addargument traxmodelpklpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument configfile defaultnone typestr requiredtrue help the config json file corresponding to the pretrained reformer model n this specifies the model architecture parser addargument pytorchdumppath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs converttraxcheckpointtopytorchargs traxmodelpklpath args configfile args pytorchdumppath 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 reformer checkpoint set parameter of one layer set torch weights for 1 to 1 comparison set torch weights for 1 to 1 comparison layernorm 1 lsh weights output intermediate weighs chunked feed forward layernorm 2 intermediate dense intermediate out reformer model word embeds output layer norm output embeddings initialise pytorch model save pytorch model required parameters
import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()


def set_param(torch_layer, weight, bias=None):
    
    assert torch_layer.weight.shape == weight.shape, f"{torch_layer} layer.weight does not match"
    torch_layer.weight = nn.Parameter(weight)
    if bias is not None:
        assert torch_layer.bias.shape == bias.shape, f"{torch_layer} layer.bias does not match"
        torch_layer.bias = nn.Parameter(bias)


def set_layer_weights_in_torch_lsh(weights, torch_layer, hidden_size):
    
    np_query_key = np.asarray(weights[0])
    np_value = np.asarray(weights[1])
    np_dense = np.asarray(weights[2])

    set_param(
        torch_layer.self_attention.query_key,
        torch.tensor(np_query_key).transpose(1, 2).contiguous().view(-1, hidden_size),
    )
    set_param(
        torch_layer.self_attention.value,
        torch.tensor(np_value).transpose(1, 2).contiguous().view(-1, hidden_size),
    )
    set_param(
        torch_layer.output.dense,
        torch.tensor(np_dense).view(-1, hidden_size).contiguous().transpose(0, 1),
    )


def set_layer_weights_in_torch_local(weights, torch_layer, hidden_size):
    
    np_query = np.asarray(weights[0])
    np_key = np.asarray(weights[1])
    np_value = np.asarray(weights[2])
    np_dense = np.asarray(weights[3])

    set_param(
        torch_layer.self_attention.query,
        torch.tensor(np_query).transpose(1, 2).contiguous().view(-1, hidden_size),
    )
    set_param(
        torch_layer.self_attention.key,
        torch.tensor(np_key).transpose(1, 2).contiguous().view(-1, hidden_size),
    )
    set_param(
        torch_layer.self_attention.value,
        torch.tensor(np_value).transpose(1, 2).contiguous().view(-1, hidden_size),
    )
    set_param(
        torch_layer.output.dense,
        torch.tensor(np_dense).view(-1, hidden_size).contiguous().transpose(0, 1),
    )


def set_block_weights_in_torch(weights, torch_block, hidden_size):
    
    layer_norm_1 = weights[0][0][0]
    layer_norm_1_weight = np.asarray(layer_norm_1[0])
    layer_norm_1_bias = np.asarray(layer_norm_1[1])
    set_param(
        torch_block.attention.layer_norm,
        torch.tensor(layer_norm_1_weight),
        torch.tensor(layer_norm_1_bias),
    )

    
    attn_weights = weights[0][1]
    if len(attn_weights) < 4:
        set_layer_weights_in_torch_lsh(attn_weights, torch_block.attention, hidden_size)
    else:
        set_layer_weights_in_torch_local(attn_weights, torch_block.attention, hidden_size)

    
    intermediate_weights = weights[2][0][1][2]

    
    if len(intermediate_weights) == 4:
        intermediate_weights = intermediate_weights[2]

    
    layer_norm_2_weight = np.asarray(intermediate_weights[0][0])
    layer_norm_2_bias = np.asarray(intermediate_weights[0][1])
    set_param(
        torch_block.feed_forward.layer_norm,
        torch.tensor(layer_norm_2_weight),
        torch.tensor(layer_norm_2_bias),
    )

    
    inter_dense_weight = np.asarray(intermediate_weights[1][0])
    inter_dense_bias = np.asarray(intermediate_weights[1][1])
    set_param(
        torch_block.feed_forward.dense.dense,
        torch.tensor(inter_dense_weight).transpose(0, 1).contiguous(),
        torch.tensor(inter_dense_bias),
    )

    
    out_dense_weight = np.asarray(intermediate_weights[4][0])
    out_dense_bias = np.asarray(intermediate_weights[4][1])
    set_param(
        torch_block.feed_forward.output.dense,
        torch.tensor(out_dense_weight).transpose(0, 1).contiguous(),
        torch.tensor(out_dense_bias),
    )


def set_model_weights_in_torch(weights, torch_model, hidden_size):
    
    torch_model_reformer = torch_model.reformer

    
    word_embeddings = np.asarray(weights[1])
    set_param(
        torch_model_reformer.embeddings.word_embeddings,
        torch.tensor(word_embeddings),
    )

    if isinstance(weights[3], tuple):
        position_embeddings = torch_model_reformer.embeddings.position_embeddings
        for emb_idx in range(len(position_embeddings.weights)):
            emb_weights = np.asarray(weights[3][emb_idx][0])
            assert (
                position_embeddings.weights[emb_idx].shape == emb_weights.shape
            ), f"{position_embeddings[emb_idx]} emb does not match"
            position_embeddings.weights[emb_idx] = nn.Parameter(torch.tensor(emb_weights))

    trax_layer_weights = weights[5]
    assert len(torch_model_reformer.encoder.layers) * 4 == len(
        trax_layer_weights
    ), "HF and trax model do not have the same number of layers"
    for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers):
        block_weights = trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
        set_block_weights_in_torch(block_weights, layer, hidden_size)

    
    layer_norm_out_weight = np.asarray(weights[7][0])
    layer_norm_out_bias = np.asarray(weights[7][1])
    set_param(
        torch_model_reformer.encoder.layer_norm,
        torch.tensor(layer_norm_out_weight),
        torch.tensor(layer_norm_out_bias),
    )

    
    output_embed_weights = np.asarray(weights[9][0])
    output_embed_bias = np.asarray(weights[9][1])
    set_param(
        torch_model.lm_head.decoder,
        torch.tensor(output_embed_weights).transpose(0, 1).contiguous(),
        torch.tensor(output_embed_bias),
    )


def convert_trax_checkpoint_to_pytorch(trax_model_pkl_path, config_file, pytorch_dump_path):
    
    config = ReformerConfig.from_json_file(config_file)
    print(f"Building PyTorch model from configuration: {config}")
    model = ReformerModelWithLMHead(config)

    with open(trax_model_pkl_path, "rb") as f:
        model_weights = pickle.load(f)["weights"]

    set_model_weights_in_torch(model_weights, model, config.hidden_size)

    
    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(
        "--trax_model_pkl_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 Reformer model. \n"
            "This specifies the model architecture."
        ),
    )
    parser.add_argument(
        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    args = parser.parse_args()
    convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)