


"""simple docstring"""







import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel




class 				UpperCamelCase_			(				pl.LightningModule):

	"""simple docstring"""





	def __init__(				self    :       Optional[Any]    ,    UpperCAmelCase__    :       Optional[Any]     ) ->			Optional[int]:
						super().__init__()
						__SCREAMING_SNAKE_CASE		     =       model
						__SCREAMING_SNAKE_CASE		     =       2
						__SCREAMING_SNAKE_CASE		     =       nn.Linear(self.model.config.hidden_size    ,    self.num_labels     )




	def 						UpperCAmelCase_   (				self    :       Tuple     ) ->			int:
						pass




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       LongformerModel.from_pretrained(lowerCAmelCase_     )
					__SCREAMING_SNAKE_CASE		     =       LightningModel(lowerCAmelCase_     )

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

					# init longformer question answering model
					__SCREAMING_SNAKE_CASE		     =       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__":
		a__       :							Optional[Any]        =  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.'''
		)
		a__       :							List[Any]        =  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"""







from abc import ABC, abstractmethod
from argparse import ArgumentParser




class 				UpperCamelCase_			(				UpperCamelCase):

	"""simple docstring"""





	@staticmethod
	@abstractmethod
	def 						UpperCAmelCase_   (				UpperCAmelCase__    :       ArgumentParser     ) ->			Optional[Any]:
						raise NotImplementedError()




	@abstractmethod
	def 						UpperCAmelCase_   (				self    :       Optional[Any]     ) ->			List[str]:
						raise NotImplementedError()
