


'''simple docstring'''



import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel





class 							__UpperCamelCase				(		pl.LightningModule					):




	def __init__(   self						,				__a							):

						'''simple docstring'''





						super().__init__()
						__a			:       Optional[Any]					    =  model
						__a			:       Any					    =  2
						__a			:       Any					    =  nn.Linear(self.model.config.hidden_size						,				self.num_labels							)





	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						pass






def   lowerCamelCase							(_SCREAMING_SNAKE_CASE      :   str						,      _SCREAMING_SNAKE_CASE      :   str						,      _SCREAMING_SNAKE_CASE      :   str   ):
					# load longformer model from model identifier
					__a			:       str					    =  LongformerModel.from_pretrained(_SCREAMING_SNAKE_CASE   )
					__a			:       Any					    =  LightningModel(_SCREAMING_SNAKE_CASE   )

					__a			:       str					    =  torch.load(_SCREAMING_SNAKE_CASE						,      map_location=torch.device('cpu'   )   )
					lightning_model.load_state_dict(ckpt['state_dict']   )

					# init longformer question answering model
					__a			:       Any					    =  LongformerForQuestionAnswering.from_pretrained(_SCREAMING_SNAKE_CASE   )

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

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


if __name__ == "__main__":
	__lowercase	:						List[str]       =	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	:						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 .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

UpperCamelCase__       :	Dict    = [
    '''python''',
    '''tqdm''',
    '''regex''',
    '''requests''',
    '''packaging''',
    '''filelock''',
    '''numpy''',
    '''tokenizers''',
    '''huggingface-hub''',
    '''safetensors''',
    '''accelerate''',
    '''pyyaml''',
]

for pkg in pkgs_to_check_at_runtime:
						if pkg in deps:
												if pkg == "tokenizers":
																		# must be loaded here, or else tqdm check may fail
																		from .utils import is_tokenizers_available

																		if not is_tokenizers_available():
																								continue  # not required, check version only if installed
												elif pkg == "accelerate":
																		# must be loaded here, or else tqdm check may fail
																		from .utils import is_accelerate_available

																		# Maybe switch to is_torch_available in the future here so that Accelerate is hard dep of
																		# Transformers with PyTorch
																		if not is_accelerate_available():
																								continue  # not required, check version only if installed

												require_version_core(deps[pkg])
						else:
												raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")


def 				lowerCAmelCase_   (	_lowerCamelCase:   List[Any]       ,      _lowerCamelCase:   Optional[int]=None					):
		require_version(deps[pkg]       ,      _lowerCamelCase					)