






import argparse

import pytorch_lightning as pl
import torch
from torch import nn

from transformers import LongformerForQuestionAnswering, LongformerModel


class      lowercase__     (			pl.LightningModule	):

	def __init__(							self		,    __UpperCAmelCase       )->					Dict:







							'''simple docstring'''





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


	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							pass







def        _a  (     UpperCamelCase_					:				str    ,   UpperCamelCase_					:				str    ,   UpperCamelCase_					:				str   ) ->   Tuple:
						"""simple docstring"""







						lowerCAmelCase__       =						LongformerModel.from_pretrained(__UpperCamelCase   )
						lowerCAmelCase__       =						LightningModel(__UpperCamelCase   )

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

						# init longformer question answering model
						lowerCAmelCase__       =						LongformerForQuestionAnswering.from_pretrained(__UpperCamelCase   )

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

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


if __name__ == "__main__":
	a_							=    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_							=    parser.parse_args()
	convert_longformer_qa_checkpoint_to_pytorch(
	    args.longformer_model, args.longformer_question_answering_ckpt_path, args.pytorch_dump_folder_path
	)














from __future__ import annotations

from collections.abc import Generator


def        _a  (     ) ->   Generator[int, None, None]:
						"""simple docstring"""







						lowerCAmelCase__       =						{}
						lowerCAmelCase__       =						2
						while True:
												lowerCAmelCase__       =						factor_map.pop(UpperCamelCase_    ,   UpperCamelCase_   )
												if factor:
																		lowerCAmelCase__       =						factor + prime
																		while x in factor_map:
																								x += factor
																		lowerCAmelCase__       =						factor
												else:
																		lowerCAmelCase__       =						prime
																		yield prime
												prime += 1


def        _a  (     UpperCamelCase_					:				float = 1e1_0   ) ->   int:
						"""simple docstring"""







						lowerCAmelCase__       =						sieve()

						lowerCAmelCase__       =						1
						while True:
												lowerCAmelCase__       =						next(UpperCamelCase_   )
												if (2 * prime * n) > limit:
																		return n
												# Ignore the next prime as the reminder will be 2.
												next(UpperCamelCase_   )
												n += 2


if __name__ == "__main__":
	print(solution())






