


import os
import sys

lowerCAmelCase__  :Optional[Any]             =							os.path.join(os.path.dirname(__file__), '''src''')
sys.path.append(SRC_DIR)


from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForMaskedLM,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    add_start_docstrings,
)


lowerCAmelCase__  :Dict             =							[
    '''torch''',
    '''numpy''',
    '''tokenizers''',
    '''filelock''',
    '''requests''',
    '''tqdm''',
    '''regex''',
    '''sentencepiece''',
    '''sacremoses''',
    '''importlib_metadata''',
    '''huggingface_hub''',
]

@add_start_docstrings(AutoConfig.__doc__	)
def 	lowerCAmelCase__			(      *a__:		int      ,  **a__:		List[Any]	)  ->		List[Any]:


							'''simple docstring'''





							return AutoConfig.from_pretrained(*a__      ,  **a__	)

@add_start_docstrings(AutoTokenizer.__doc__	)
def 	lowerCAmelCase__			(      *a__:		List[str]      ,  **a__:		Optional[Any]	)  ->		List[Any]:


							'''simple docstring'''





							return AutoTokenizer.from_pretrained(*a__      ,  **a__	)

@add_start_docstrings(AutoModel.__doc__	)
def 	lowerCAmelCase__			(      *a__:		Tuple      ,  **a__:		Dict	)  ->		Optional[int]:


							'''simple docstring'''





							return AutoModel.from_pretrained(*a__      ,  **a__	)

@add_start_docstrings(AutoModelForCausalLM.__doc__	)
def 	lowerCAmelCase__			(      *a__:		Union[str, Any]      ,  **a__:		List[Any]	)  ->		Union[str, Any]:


							'''simple docstring'''





							return AutoModelForCausalLM.from_pretrained(*a__      ,  **a__	)

@add_start_docstrings(AutoModelForMaskedLM.__doc__	)
def 	lowerCAmelCase__			(      *a__:		Any      ,  **a__:		List[str]	)  ->		Optional[Any]:


							'''simple docstring'''





							return AutoModelForMaskedLM.from_pretrained(*a__      ,  **a__	)

@add_start_docstrings(AutoModelForSequenceClassification.__doc__	)
def 	lowerCAmelCase__			(      *a__:		Dict      ,  **a__:		str	)  ->		int:


							'''simple docstring'''





							return AutoModelForSequenceClassification.from_pretrained(*a__      ,  **a__	)







@add_start_docstrings(AutoModelForQuestionAnswering.__doc__	)
def 	lowerCAmelCase__			(      *a__:		Union[str, Any]      ,  **a__:		List[Any]	)  ->		List[str]:


							'''simple docstring'''





							return AutoModelForQuestionAnswering.from_pretrained(*a__      ,  **a__	)




import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel

class 		__a				(    pl.LightningModule      ):
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE						)			->	List[Any]:




									"""simple docstring"""





									super().__init__()
									_UpperCAmelCase					=      model
									_UpperCAmelCase					=      2
									_UpperCAmelCase					=      nn.Linear(self.model.config.hidden_size      ,   self.num_labels						)






		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									pass



def 	lowerCAmelCase__			(      a__:		str      ,  a__:		str      ,  a__:		str	)  ->		Tuple:


							'''simple docstring'''





							_UpperCAmelCase					=      LongformerModel.from_pretrained(a__	)
							_UpperCAmelCase					=      LightningModel(a__	)

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

							# init longformer question answering model
							_UpperCAmelCase					=      LongformerForQuestionAnswering.from_pretrained(a__	)

							# 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(a__	)

							print(F'''Conversion successful. Model saved under {pytorch_dump_folder_path}'''	)


if __name__ == "__main__":
		lowerCAmelCase__  :Optional[int]             =							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.'''
		)
		lowerCAmelCase__  :Dict             =							parser.parse_args()
		convert_longformer_qa_checkpoint_to_pytorch(
		    args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
		)
