


'''simple docstring'''






import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel






class    __magic_name__				(						pl.LightningModule							):






       def __init__(     self    ,				snake_case)		->    List[Any]:




          '''simple docstring'''
          super().__init__()
          _UpperCAmelCase  :  Optional[Any]				       =model
          _UpperCAmelCase  :  Union[str, Any]				       =2
          _UpperCAmelCase  :  Optional[int]				       =nn.Linear(self.model.config.hidden_size    ,				self.num_labels)






       def      lowerCAmelCase						(     self)		->    Tuple:




          '''simple docstring'''
          pass






def lowerCamelCase__     (		__lowerCamelCase			:				str			,       __lowerCamelCase			:				str			,       __lowerCamelCase			:				str				):







   '''simple docstring'''




   _UpperCAmelCase  :  List[Any]				       =LongformerModel.from_pretrained(__lowerCamelCase				)
   _UpperCAmelCase  :  Optional[int]				       =LightningModel(__lowerCamelCase				)

   _UpperCAmelCase  :  Any				       =torch.load(__lowerCamelCase			,       map_location=torch.device('cpu'				)				)
   lightning_model.load_state_dict(ckpt['state_dict']				)

   # init longformer question answering model
   _UpperCAmelCase  :  Tuple				       =LongformerForQuestionAnswering.from_pretrained(__lowerCamelCase				)

   # transfer weights
   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()

   # save model
   longformer_for_qa.save_pretrained(__lowerCamelCase				)

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


if __name__ == "__main__":
 lowercase				=argparse.ArgumentParser()
 # Required parameters
 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.'
 )
 lowercase				=parser.parse_args()
 convert_longformer_qa_checkpoint_to_pytorch(
     args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
 )





'''simple docstring'''






import warnings
from functools import wraps
from typing import Callable





def lowerCamelCase__     (		__lowerCamelCase			:				Callable				):







   '''simple docstring'''




   @wraps(__lowerCamelCase				)
   def _inner_fn(*__lowerCamelCase			:				int			,       **__lowerCamelCase			:				Optional[Any]				):
      warnings.warn(
          (f"'{fn.__name__}' is experimental and might be subject to breaking changes in the future.")			,       __lowerCamelCase			,       )
      return fn(*__lowerCamelCase			,       **__lowerCamelCase				)

   return _inner_fn

