codingutf8 2018 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 roberta checkpoint import argparse import pytorchlightning as pl import torch from torch import nn from transformers import longformerforquestionanswering longformermodel class lightningmodelpl lightningmodule def initself model super init self model model self numlabels 2 self qaoutputs nn linearself model config hiddensize self numlabels implement only because lightning requires to do so def forwardself pass def convertlongformerqacheckpointtopytorch longformermodel str longformerquestionansweringckptpath str pytorchdumpfolderpath str load longformer model from model identifier longformer longformermodel frompretrainedlongformermodel lightningmodel lightningmodellongformer ckpt torch loadlongformerquestionansweringckptpath maplocationtorch devicecpu lightningmodel loadstatedictckptstatedict init longformer question answering model longformerforqa longformerforquestionanswering frompretrainedlongformermodel transfer weights longformerforqa longformer loadstatedictlightningmodel model statedict longformerforqa qaoutputs loadstatedictlightningmodel qaoutputs statedict longformerforqa eval save model longformerforqa savepretrainedpytorchdumpfolderpath printfconversion successful model saved under pytorchdumpfolderpath if name main parser argparse argumentparser required parameters parser addargument longformermodel defaultnone typestr requiredtrue helpmodel identifier of longformer should be either longformerbase4096 or longformerlarge4096 parser addargument longformerquestionansweringckptpath defaultnone typestr requiredtrue helppath the official pytorch lightning checkpoint parser addargument pytorchdumpfolderpath defaultnone typestr requiredtrue helppath to the output pytorch model args parser parseargs convertlongformerqacheckpointtopytorch args longformermodel args longformerquestionansweringckptpath args pytorchdumpfolderpath coding utf 8 2018 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 roberta checkpoint implement only because lightning requires to do so load longformer model from model identifier init longformer question answering model transfer weights save model required parameters
import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel


class LightningModel(pl.LightningModule):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.num_labels = 2
        self.qa_outputs = nn.Linear(self.model.config.hidden_size, self.num_labels)

    
    def forward(self):
        pass


def convert_longformer_qa_checkpoint_to_pytorch(
    longformer_model: str, longformer_question_answering_ckpt_path: str, pytorch_dump_folder_path: str
):
    
    longformer = LongformerModel.from_pretrained(longformer_model)
    lightning_model = LightningModel(longformer)

    ckpt = torch.load(longformer_question_answering_ckpt_path, map_location=torch.device("cpu"))
    lightning_model.load_state_dict(ckpt["state_dict"])

    
    longformer_for_qa = LongformerForQuestionAnswering.from_pretrained(longformer_model)

    
    longformer_for_qa.longformer.load_state_dict(lightning_model.model.state_dict())
    longformer_for_qa.qa_outputs.load_state_dict(lightning_model.qa_outputs.state_dict())
    longformer_for_qa.eval()

    
    longformer_for_qa.save_pretrained(pytorch_dump_folder_path)

    print(f"Conversion successful. Model saved under {pytorch_dump_folder_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        "--longformer_model",
        default=None,
        type=str,
        required=True,
        help="model identifier of longformer. Should be either `longformer-base-4096` or `longformer-large-4096`.",
    )
    parser.add_argument(
        "--longformer_question_answering_ckpt_path",
        default=None,
        type=str,
        required=True,
        help="Path the official PyTorch Lightning Checkpoint.",
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    args = parser.parse_args()
    convert_longformer_qa_checkpoint_to_pytorch(
        args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
    )