


from __future__ import annotations

def 	lowerCAmelCase__			(      a__:		Union[str, Any]   ,  a__:		Any	)  ->		list[str]:


							'''simple docstring'''





							if partitions <= 0:
														raise ValueError('partitions must be a positive number!'	)
							if partitions > number_of_bytes:
														raise ValueError('partitions can not > number_of_bytes!'	)
							_UpperCAmelCase					=      number_of_bytes // partitions
							_UpperCAmelCase					=      []
							for i in range(_lowerCamelCase	):
														_UpperCAmelCase					=      i * bytes_per_partition + 1
														_UpperCAmelCase					=      (
														    number_of_bytes if i == partitions - 1 else (i + 1) * bytes_per_partition
														)
														allocation_list.append(F'''{start_bytes}-{end_bytes}'''	)
							return allocation_list


if __name__ == "__main__":
		import doctest

		doctest.testmod()




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
		)
