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 fnet checkpoint import argparse import torch from flax training checkpoints import restorecheckpoint from transformers import fnetconfig fnetforpretraining from transformers utils import logging logging setverbosityinfo def convertflaxcheckpointtopytorchflaxcheckpointpath fnetconfigfile savepath initialise pytorch model config fnetconfig fromjsonfilefnetconfigfile printfbuilding pytorch model from configuration config fnetpretrainingmodel fnetforpretrainingconfig checkpointdict restorecheckpointflaxcheckpointpath none pretrainedmodelparams checkpointdicttarget embeddings position ids statedict fnetpretrainingmodel statedict positionids statedictfnet embeddings positionids newstatedict fnet embeddings positionids positionids embedding layers newstatedictfnet embeddings wordembeddings weight torch tensor pretrainedmodelparamsencoderembedderwordembedding newstatedictfnet embeddings positionembeddings weight torch tensor pretrainedmodelparamsencoderembedderpositionembedding0 newstatedictfnet embeddings tokentypeembeddings weight torch tensor pretrainedmodelparamsencoderembeddertypeembedding newstatedictfnet embeddings projection weight torch tensor pretrainedmodelparamsencoderembedderhiddenmappinginkernel t newstatedictfnet embeddings projection bias torch tensor pretrainedmodelparamsencoderembedderhiddenmappinginbias newstatedictfnet embeddings layernorm weight torch tensor pretrainedmodelparamsencoderembedderlayernormscale newstatedictfnet embeddings layernorm bias torch tensor pretrainedmodelparamsencoderembedderlayernormbias encoder layers for layer in rangeconfig numhiddenlayers newstatedictffnet encoder layer layer fourier output layernorm weight torch tensor pretrainedmodelparamsencoderfencoderlayermixinglayernormscale newstatedictffnet encoder layer layer fourier output layernorm bias torch tensor pretrainedmodelparamsencoderfencoderlayermixinglayernormbias newstatedictffnet encoder layer layer intermediate dense weight torch tensor pretrainedmodelparamsencoderffeedforwardlayerintermediatekernel t newstatedictffnet encoder layer layer intermediate dense bias torch tensor pretrainedmodelparamsencoderffeedforwardlayerintermediatebias newstatedictffnet encoder layer layer output dense weight torch tensor pretrainedmodelparamsencoderffeedforwardlayeroutputkernel t newstatedictffnet encoder layer layer output dense bias torch tensor pretrainedmodelparamsencoderffeedforwardlayeroutputbias newstatedictffnet encoder layer layer output layernorm weight torch tensor pretrainedmodelparamsencoderfencoderlayeroutputlayernormscale newstatedictffnet encoder layer layer output layernorm bias torch tensor pretrainedmodelparamsencoderfencoderlayeroutputlayernormbias pooler layers newstatedictfnet pooler dense weight torch tensorpretrainedmodelparamsencoderpoolerkernel t newstatedictfnet pooler dense bias torch tensorpretrainedmodelparamsencoderpoolerbias masked lm layers newstatedictcls predictions transform dense weight torch tensor pretrainedmodelparamspredictionsdensekernel t newstatedictcls predictions transform dense bias torch tensor pretrainedmodelparamspredictionsdensebias newstatedictcls predictions transform layernorm weight torch tensor pretrainedmodelparamspredictionslayernormscale newstatedictcls predictions transform layernorm bias torch tensor pretrainedmodelparamspredictionslayernormbias newstatedictcls predictions decoder weight torch tensor pretrainedmodelparamsencoderembedderwordembedding newstatedictcls predictions decoder bias torch tensor pretrainedmodelparamspredictionsoutputoutputbias newstatedictcls predictions bias torch tensorpretrainedmodelparamspredictionsoutputoutputbias seq relationship layers newstatedictcls seqrelationship weight torch tensor pretrainedmodelparamsclassificationoutputkernel newstatedictcls seqrelationship bias torch tensor pretrainedmodelparamsclassificationoutputbias load state dict fnetpretrainingmodel loadstatedictnewstatedict save pretrained printfsaving pretrained model to savepath fnetpretrainingmodel savepretrainedsavepath if name main parser argparse argumentparser required parameters parser addargument flaxcheckpointpath defaultnone typestr requiredtrue helppath to the tensorflow checkpoint path parser addargument fnetconfigfile defaultnone typestr requiredtrue help the config json file corresponding to the pretrained fnet model n this specifies the model architecture parser addargumentsavepath defaultnone typestr requiredtrue helppath to the output model args parser parseargs convertflaxcheckpointtopytorchargs flaxcheckpointpath args fnetconfigfile args savepath 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 fnet checkpoint initialise pytorch model embeddings position ids embedding layers encoder layers pooler layers masked lm layers seq relationship layers load state dict save pretrained required parameters
import argparse

import torch
from flax.training.checkpoints import restore_checkpoint

from transformers import FNetConfig, FNetForPreTraining
from transformers.utils import logging


logging.set_verbosity_info()


def convert_flax_checkpoint_to_pytorch(flax_checkpoint_path, fnet_config_file, save_path):
    
    config = FNetConfig.from_json_file(fnet_config_file)
    print(f"Building PyTorch model from configuration: {config}")
    fnet_pretraining_model = FNetForPreTraining(config)

    checkpoint_dict = restore_checkpoint(flax_checkpoint_path, None)
    pretrained_model_params = checkpoint_dict["target"]

    
    
    state_dict = fnet_pretraining_model.state_dict()

    position_ids = state_dict["fnet.embeddings.position_ids"]
    new_state_dict = {"fnet.embeddings.position_ids": position_ids}
    
    new_state_dict["fnet.embeddings.word_embeddings.weight"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]
    )
    new_state_dict["fnet.embeddings.position_embeddings.weight"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["position"]["embedding"][0]
    )
    new_state_dict["fnet.embeddings.token_type_embeddings.weight"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["type"]["embedding"]
    )
    new_state_dict["fnet.embeddings.projection.weight"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["kernel"]
    ).T
    new_state_dict["fnet.embeddings.projection.bias"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["bias"]
    )
    new_state_dict["fnet.embeddings.LayerNorm.weight"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["layer_norm"]["scale"]
    )
    new_state_dict["fnet.embeddings.LayerNorm.bias"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["layer_norm"]["bias"]
    )

    
    for layer in range(config.num_hidden_layers):
        new_state_dict[f"fnet.encoder.layer.{layer}.fourier.output.LayerNorm.weight"] = torch.tensor(
            pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["scale"]
        )
        new_state_dict[f"fnet.encoder.layer.{layer}.fourier.output.LayerNorm.bias"] = torch.tensor(
            pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["bias"]
        )

        new_state_dict[f"fnet.encoder.layer.{layer}.intermediate.dense.weight"] = torch.tensor(
            pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["kernel"]
        ).T
        new_state_dict[f"fnet.encoder.layer.{layer}.intermediate.dense.bias"] = torch.tensor(
            pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["bias"]
        )

        new_state_dict[f"fnet.encoder.layer.{layer}.output.dense.weight"] = torch.tensor(
            pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["kernel"]
        ).T
        new_state_dict[f"fnet.encoder.layer.{layer}.output.dense.bias"] = torch.tensor(
            pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["bias"]
        )

        new_state_dict[f"fnet.encoder.layer.{layer}.output.LayerNorm.weight"] = torch.tensor(
            pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["scale"]
        )
        new_state_dict[f"fnet.encoder.layer.{layer}.output.LayerNorm.bias"] = torch.tensor(
            pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["bias"]
        )

    
    new_state_dict["fnet.pooler.dense.weight"] = torch.tensor(pretrained_model_params["encoder"]["pooler"]["kernel"]).T
    new_state_dict["fnet.pooler.dense.bias"] = torch.tensor(pretrained_model_params["encoder"]["pooler"]["bias"])

    
    new_state_dict["cls.predictions.transform.dense.weight"] = torch.tensor(
        pretrained_model_params["predictions_dense"]["kernel"]
    ).T
    new_state_dict["cls.predictions.transform.dense.bias"] = torch.tensor(
        pretrained_model_params["predictions_dense"]["bias"]
    )
    new_state_dict["cls.predictions.transform.LayerNorm.weight"] = torch.tensor(
        pretrained_model_params["predictions_layer_norm"]["scale"]
    )
    new_state_dict["cls.predictions.transform.LayerNorm.bias"] = torch.tensor(
        pretrained_model_params["predictions_layer_norm"]["bias"]
    )
    new_state_dict["cls.predictions.decoder.weight"] = torch.tensor(
        pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]
    )
    new_state_dict["cls.predictions.decoder.bias"] = torch.tensor(
        pretrained_model_params["predictions_output"]["output_bias"]
    )
    new_state_dict["cls.predictions.bias"] = torch.tensor(pretrained_model_params["predictions_output"]["output_bias"])

    
    new_state_dict["cls.seq_relationship.weight"] = torch.tensor(
        pretrained_model_params["classification"]["output_kernel"]
    )
    new_state_dict["cls.seq_relationship.bias"] = torch.tensor(
        pretrained_model_params["classification"]["output_bias"]
    )

    
    fnet_pretraining_model.load_state_dict(new_state_dict)

    
    print(f"Saving pretrained model to {save_path}")
    fnet_pretraining_model.save_pretrained(save_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--flax_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
    )
    parser.add_argument(
        "--fnet_config_file",
        default=None,
        type=str,
        required=True,
        help=(
            "The config json file corresponding to the pre-trained FNet model. \n"
            "This specifies the model architecture."
        ),
    )
    parser.add_argument("--save_path", default=None, type=str, required=True, help="Path to the output model.")
    args = parser.parse_args()
    convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)